<commit_before>extern crate eventual;
extern crate sbahn;

use eventual::*;
use sbahn::client;
use sbahn::handler;
use sbahn::message::*;
use sbahn::storage::HashMapBackend;
use sbahn::storage_node::StorageNode;
use std::net::{Ipv4Addr, SocketAddrV4, TcpListener};
use std::thread;
use std::time::Duration;

// Milis to wait before trying to connect to any node.
static DELAY: u64 = 50;

static mut PORT: u16 = 1100;
/// Obtain an open port
fn get_port() -> u16 {
    let mut port = 0;
    loop {
        unsafe {
          PORT += 1;
          port = PORT;
        }
        let addr = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), port);
        if TcpListener::bind(&addr).is_ok() {
            // Check wether the port is open, and only return it if it is.
            return port;
        }
    }
}

fn get_storage_node(pos: usize, shard_count: usize) -> SocketAddrV4 {
    let addr = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), get_port());
    let mut sn: StorageNode<HashMapBackend> = StorageNode::new(&addr, pos, shard_count);
    thread::spawn(move || {
        &sn.listen();
    });
    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening
    addr
}

fn setup_handler_node(shards: &Vec<Vec<SocketAddrV4>>) -> SocketAddrV4 {
    let shards = shards.clone();
    let addr = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), get_port());
    thread::spawn(move || {
        let _ = handler::listen(&addr, &shards);
    });
    thread::sleep(Duration::from_millis(DELAY));  // Wait for handler node to start listening
    addr
}

/// Listen on `address` for incoming client requests, and do nothing.
pub fn dead_node(address: &SocketAddrV4) -> Future<(), ()> {
    let address = address.to_owned();

    let read_timeout = Some(Duration::from_millis(300));
    let write_timeout = Some(Duration::from_millis(300));

    Future::spawn(move || {
        match TcpListener::bind(&address) {
            Ok(listener) => {
                // Accept connections and process them, spawning a new thread for each one.
                for stream in listener.incoming() {
                    match stream {
                        Ok(stream) => {
                            let _ = stream.set_read_timeout(read_timeout);
                            let _ = stream.set_write_timeout(write_timeout);
                            thread::spawn(|| {
                                thread::sleep(Duration::from_millis(500));
                            });
                        }
                        Err(e) => panic!("Connection failed!: {:?}", e),
                    }
                }
            }
            Err(e) => panic!("Error while binding to {:?}: {:?}", address, e),
        }
    })
}

fn get_dead_storage_node() -> SocketAddrV4 {
    let addr = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), get_port());
    thread::spawn(move || {
        let _ = dead_node(&addr);
    });
    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening
    addr
}

fn setup_cluster() -> (SocketAddrV4, Vec<Vec<SocketAddrV4>>) {
    setup_bad_cluster(0)
}

fn setup_bad_cluster(bad_nodes: u32) -> (SocketAddrV4, Vec<Vec<SocketAddrV4>>) {
    let mut shards: Vec<Vec<SocketAddrV4>> = vec![];
    for i in 0..3 {
        let mut shard: Vec<SocketAddrV4> = vec![];
        for _ in 0..bad_nodes {
            shard.push(get_dead_storage_node());
        }
        for _ in bad_nodes..3 {
            shard.push(get_storage_node(i, 3));
        }
        shards.push(shard);
    }

    (setup_handler_node(&shards), shards)
}

fn write_to_storage_node(target: &SocketAddrV4, key: &Key, value: &Vec<u8>, timestamp: u64) {
    let request = InternodeRequest::Write {
        key: key.to_owned(),
        value: Value::Value {
            content: value.to_owned(),
            timestamp: timestamp,
        },
    };
    let r: Future<InternodeResponse, Error> = client::Client::send_to_node(target, &request);
    let _ = r.await();
}

fn key_and_value() -> (Key, Vec<u8>) {
    let key = Key {
        dataset: vec![1, 2, 3],
        pkey: vec![4, 5, 6],
        lkey: vec![7, 8, 9],
    };
    let value: Vec<u8> = vec![9, 8, 7];
    (key, value)
}

#[test]
fn read_what_you_insert() {
    let (handler_addr, _) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    let client = client::Client::new(vec![handler_addr]);
    // Should succeed
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::WriteAck {key, ..} => assert_eq!(key, local_key),
        _ => assert!(false),
    }
    // Should succeed
    let r = client.get(&local_key);
    let r = r.await().unwrap();
    match r.message {
        Response::Value {key, value} => {
            assert_eq!(key, local_key);
            match value {
                Value::Value {content, ..} => assert_eq!(content, local_value),
                _ => assert!(false),
            }
        },
        _ => assert!(false),
    }
}

#[test]
fn read_consistency_one_all_nodes_available() {
    let (handler_addr, shards) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    for shard in shards {
        for node in shard {
            write_to_storage_node(&node, &local_key, &local_value, 100000);
        }
    }
    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::One);

    // Should succeed
    let r = client.get(&local_key);
    let r = r.await().unwrap();
    match r.message {
        Response::Value {key, value} => {
            assert_eq!(key, local_key);
            match value {
                Value::Value {content, ..} => assert_eq!(content, local_value),
                _ => assert!(false),
            }
        },
        _ => assert!(false),
    }
}

#[test]
fn read_consistency_one_one_node_available() {
    let (handler_addr, shards) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    // Write to only one node (local_key corresponds to shard 2).
    write_to_storage_node(&shards[2][0], &local_key, &local_value, 100000);

    thread::sleep(Duration::from_millis(DELAY*3));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::One);

    // Should succeed
    let r = client.get(&local_key);
    let r = r.await().unwrap();
    match r.message {
        Response::Value {key, value} => {
            assert_eq!(key, local_key);
            match value {
                Value::Value {content, ..} => assert_eq!(content, local_value),
                _ => assert!(false),
            }
        },
        _ => assert!(false),
    }
}

#[test]
fn read_consistency_latest_all_same() {
    // Should succeed
    let (handler_addr, shards) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    for shard in shards {
        for node in shard {
            write_to_storage_node(&node, &local_key, &local_value, 100000);
        }
    }
    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should succeed
    let r = client.get(&local_key);
    let r = r.await().unwrap();
    match r.message {
        Response::Value {key, value} => {
            assert_eq!(key, local_key);
            match value {
                Value::Value {content, ..} => assert_eq!(content, local_value),
                _ => assert!(false),
            }
        },
        _ => assert!(false),
    }
}

#[test]
fn read_consistency_latest_all_different() {
    // Should succeed
}

#[test]
fn read_consistency_latest_one_node_available() {
    // Should fail
    let (handler_addr, shards) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    // Write to only one node (local_key corresponds to shard 2).
    write_to_storage_node(&shards[2][0], &local_key, &local_value, 100000);

    thread::sleep(Duration::from_millis(DELAY*3));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should succeed
    let r = client.get(&local_key);
    let r = r.await().unwrap();
    match r.message {
        Response::Error {key, ..} => {
            assert_eq!(key, local_key);
        }
        _ => assert!(false),
    }
}

#[test]
fn write_consistency_one_all_available() {
    // Should succeed
    let (handler_addr, _) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should succeed
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::WriteAck {key, timestamp} => {
            assert_eq!(key, local_key);
            assert!(timestamp > 0);
        },
        _ => assert!(false),
    }
}

#[test]
fn write_consistency_one_none_available() {
    // Should fail
    let (handler_addr, _) = setup_bad_cluster(3);
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::One);

    // Should succeed
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::Error {key, message} => {
            assert_eq!(key, local_key);
            assert_eq!("Quorum write could not be accomplished.".to_string(), message);
        },
        _ => assert!(false),
    }
}

#[test]
fn write_consistency_latest_all_available() {
    let (handler_addr, _) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should succeed
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::WriteAck {key, timestamp} => {
            assert_eq!(key, local_key);
            assert!(timestamp > 0);
        },
        _ => assert!(false),
    }
}

#[test]
fn write_consistency_latest_quorum_available() {
    let (handler_addr, _) = setup_bad_cluster(1);
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should succeed
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::WriteAck {key, timestamp} => {
            assert_eq!(key, local_key);
            assert!(timestamp > 0);
        },
        _ => assert!(false),
    }
}

#[test]
fn write_consistency_latest_one_available() {
    let (handler_addr, _) = setup_bad_cluster(2);
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should fail
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::Error {key, message} => {
            assert_eq!(local_key, key);
            assert_eq!("Quorum write could not be accomplished.".to_string(), message);
        }
        _ => assert!(false),
    }
}

#[test]
fn write_consistency_latest_none_available() {
    // Should fail
    let (handler_addr, _) = setup_bad_cluster(3);
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should fail
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::Error {key, message} => {
            assert_eq!(local_key, key);
            assert_eq!("Quorum write could not be accomplished.".to_string(), message);
        }
        _ => assert!(false),
    }
}

#[test]
fn single_node() {
    let addr = get_storage_node(0, 1);
    let (insert_key, _) = key_and_value();

    {
        let content = InternodeRequest::Write {
            key: insert_key.to_owned(),
            value: Value::Value {
                content: vec![1],
                timestamp: 10000000,
            },
        };
        let addr = &addr.to_owned();
        let r: Future<InternodeResponse, Error> = client::Client::send_to_node(addr, &content);
        match r.await().unwrap() {
            InternodeResponse::WriteAck {key, timestamp} => {
                assert_eq!(key, insert_key);
                assert_eq!(timestamp, 10000000);
            },
            e => panic!("{:?}", e),
        }
    }
    {
        let content = InternodeRequest::Read {
            key: insert_key.to_owned(),
        };
        let addr = &addr.to_owned();
        let r: Future<InternodeResponse, Error> = client::Client::send_to_node(addr, &content);
        let r = r.await();
        match r {
            Ok(r) => match r {
                InternodeResponse::Value {key, value} => {
                    assert_eq!(key, insert_key);
                    match value {
                        Value::Value {content, timestamp} => {
                            assert_eq!(&content[..], &[1][..]);
                            assert_eq!(timestamp, 10000000);
                        },
                        _ => panic!(),
                    }
                },
                _ => panic!(),
            },
            Err(_) => panic!(),
        }
    }
}
<commit_msg>Travis CI is failing, but it works locally :-/<commit_after>extern crate eventual;
extern crate sbahn;

use eventual::*;
use sbahn::client;
use sbahn::handler;
use sbahn::message::*;
use sbahn::storage::HashMapBackend;
use sbahn::storage_node::StorageNode;
use std::net::{Ipv4Addr, SocketAddrV4, TcpListener};
use std::thread;
use std::time::Duration;

// Milis to wait before trying to connect to any node.
static DELAY: u64 = 100;

static mut PORT: u16 = 1100;
/// Obtain an open port
fn get_port() -> u16 {
    let mut port = 0;
    loop {
        unsafe {
          PORT += 1;
          port = PORT;
        }
        let addr = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), port);
        if TcpListener::bind(&addr).is_ok() {
            // Check wether the port is open, and only return it if it is.
            return port;
        }
    }
}

fn get_storage_node(pos: usize, shard_count: usize) -> SocketAddrV4 {
    let addr = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), get_port());
    let mut sn: StorageNode<HashMapBackend> = StorageNode::new(&addr, pos, shard_count);
    thread::spawn(move || {
        &sn.listen();
    });
    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening
    addr
}

fn setup_handler_node(shards: &Vec<Vec<SocketAddrV4>>) -> SocketAddrV4 {
    let shards = shards.clone();
    let addr = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), get_port());
    thread::spawn(move || {
        let _ = handler::listen(&addr, &shards);
    });
    thread::sleep(Duration::from_millis(DELAY));  // Wait for handler node to start listening
    addr
}

/// Listen on `address` for incoming client requests, and do nothing.
pub fn dead_node(address: &SocketAddrV4) -> Future<(), ()> {
    let address = address.to_owned();

    let read_timeout = Some(Duration::from_millis(300));
    let write_timeout = Some(Duration::from_millis(300));

    Future::spawn(move || {
        match TcpListener::bind(&address) {
            Ok(listener) => {
                // Accept connections and process them, spawning a new thread for each one.
                for stream in listener.incoming() {
                    match stream {
                        Ok(stream) => {
                            let _ = stream.set_read_timeout(read_timeout);
                            let _ = stream.set_write_timeout(write_timeout);
                            thread::spawn(|| {
                                thread::sleep(Duration::from_millis(500));
                            });
                        }
                        Err(e) => panic!("Connection failed!: {:?}", e),
                    }
                }
            }
            Err(e) => panic!("Error while binding to {:?}: {:?}", address, e),
        }
    })
}

fn get_dead_storage_node() -> SocketAddrV4 {
    let addr = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), get_port());
    thread::spawn(move || {
        let _ = dead_node(&addr);
    });
    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening
    addr
}

fn setup_cluster() -> (SocketAddrV4, Vec<Vec<SocketAddrV4>>) {
    setup_bad_cluster(0)
}

fn setup_bad_cluster(bad_nodes: u32) -> (SocketAddrV4, Vec<Vec<SocketAddrV4>>) {
    let mut shards: Vec<Vec<SocketAddrV4>> = vec![];
    for i in 0..3 {
        let mut shard: Vec<SocketAddrV4> = vec![];
        for _ in 0..bad_nodes {
            shard.push(get_dead_storage_node());
        }
        for _ in bad_nodes..3 {
            shard.push(get_storage_node(i, 3));
        }
        shards.push(shard);
    }

    (setup_handler_node(&shards), shards)
}

fn write_to_storage_node(target: &SocketAddrV4, key: &Key, value: &Vec<u8>, timestamp: u64) {
    let request = InternodeRequest::Write {
        key: key.to_owned(),
        value: Value::Value {
            content: value.to_owned(),
            timestamp: timestamp,
        },
    };
    let r: Future<InternodeResponse, Error> = client::Client::send_to_node(target, &request);
    let _ = r.await();
}

fn key_and_value() -> (Key, Vec<u8>) {
    let key = Key {
        dataset: vec![1, 2, 3],
        pkey: vec![4, 5, 6],
        lkey: vec![7, 8, 9],
    };
    let value: Vec<u8> = vec![9, 8, 7];
    (key, value)
}

#[test]
fn read_what_you_insert() {
    let (handler_addr, _) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    let client = client::Client::new(vec![handler_addr]);
    // Should succeed
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::WriteAck {key, ..} => assert_eq!(key, local_key),
        _ => assert!(false),
    }
    // Should succeed
    let r = client.get(&local_key);
    let r = r.await().unwrap();
    match r.message {
        Response::Value {key, value} => {
            assert_eq!(key, local_key);
            match value {
                Value::Value {content, ..} => assert_eq!(content, local_value),
                _ => assert!(false),
            }
        },
        _ => assert!(false),
    }
}

#[test]
fn read_consistency_one_all_nodes_available() {
    let (handler_addr, shards) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    for shard in shards {
        for node in shard {
            write_to_storage_node(&node, &local_key, &local_value, 100000);
        }
    }
    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::One);

    // Should succeed
    let r = client.get(&local_key);
    let r = r.await().unwrap();
    match r.message {
        Response::Value {key, value} => {
            assert_eq!(key, local_key);
            match value {
                Value::Value {content, ..} => assert_eq!(content, local_value),
                _ => assert!(false),
            }
        },
        _ => assert!(false),
    }
}

#[test]
fn read_consistency_one_one_node_available() {
    let (handler_addr, shards) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    // Write to only one node (local_key corresponds to shard 2).
    write_to_storage_node(&shards[2][0], &local_key, &local_value, 100000);

    thread::sleep(Duration::from_millis(DELAY*3));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::One);

    // Should succeed
    let r = client.get(&local_key);
    let r = r.await().unwrap();
    match r.message {
        Response::Value {key, value} => {
            assert_eq!(key, local_key);
            match value {
                Value::Value {content, ..} => assert_eq!(content, local_value),
                _ => assert!(false),
            }
        },
        _ => assert!(false),
    }
}

#[test]
fn read_consistency_latest_all_same() {
    // Should succeed
    let (handler_addr, shards) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    for shard in shards {
        for node in shard {
            write_to_storage_node(&node, &local_key, &local_value, 100000);
        }
    }
    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should succeed
    let r = client.get(&local_key);
    let r = r.await().unwrap();
    match r.message {
        Response::Value {key, value} => {
            assert_eq!(key, local_key);
            match value {
                Value::Value {content, ..} => assert_eq!(content, local_value),
                _ => assert!(false),
            }
        },
        _ => assert!(false),
    }
}

#[test]
fn read_consistency_latest_all_different() {
    // Should succeed
}

#[test]
fn read_consistency_latest_one_node_available() {
    let (handler_addr, shards) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    // Write to only one node (local_key corresponds to shard 2).
    write_to_storage_node(&shards[2][0], &local_key, &local_value, 100000);

    thread::sleep(Duration::from_millis(DELAY*3));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should fail
    let r = client.get(&local_key);
    let r = r.await().unwrap();
    match r.message {
        Response::Error {key, ..} => {
            assert_eq!(key, local_key);
        }
        _ => assert!(false),
    }
}

#[test]
fn write_consistency_one_all_available() {
    let (handler_addr, _) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should succeed
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::WriteAck {key, timestamp} => {
            assert_eq!(key, local_key);
            assert!(timestamp > 0);
        },
        _ => assert!(false),
    }
}

#[test]
fn write_consistency_one_none_available() {
    let (handler_addr, _) = setup_bad_cluster(3);
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::One);

    // Should fail
    match client.insert(&local_key, &local_value).await() {
        Ok(r) => {
            match r.message {
                Response::Error {key, message} => {
                    assert_eq!(key, local_key);
                    assert_eq!("Quorum write could not be accomplished.".to_string(), message);
                },
                _ => assert!(false),
            }
        }
        Err(_) => assert!(true),
    }
}

#[test]
fn write_consistency_latest_all_available() {
    let (handler_addr, _) = setup_cluster();
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should succeed
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::WriteAck {key, timestamp} => {
            assert_eq!(key, local_key);
            assert!(timestamp > 0);
        },
        _ => assert!(false),
    }
}

#[test]
fn write_consistency_latest_quorum_available() {
    let (handler_addr, _) = setup_bad_cluster(1);
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should succeed
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::WriteAck {key, timestamp} => {
            assert_eq!(key, local_key);
            assert!(timestamp > 0);
        },
        _ => assert!(false),
    }
}

#[test]
fn write_consistency_latest_one_available() {
    let (handler_addr, _) = setup_bad_cluster(2);
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should fail
    let r = client.insert(&local_key, &local_value);
    let r = r.await().unwrap();
    match r.message {
        Response::Error {key, message} => {
            assert_eq!(local_key, key);
            assert_eq!("Quorum write could not be accomplished.".to_string(), message);
        }
        _ => assert!(false),
    }
}

#[test]
fn write_consistency_latest_none_available() {
    // Should fail
    let (handler_addr, _) = setup_bad_cluster(3);
    let (local_key, local_value) = key_and_value();

    thread::sleep(Duration::from_millis(DELAY));  // Wait for storage node to start listening

    let client = client::Client::with_consistency(vec![handler_addr], Consistency::Latest);

    // Should fail
    match client.insert(&local_key, &local_value).await() {
        Ok(r) => {
            match r.message {
                Response::Error {key, message} => {
                    assert_eq!(local_key, key);
                    assert_eq!("Quorum write could not be accomplished.".to_string(), message);
                }
                _ => assert!(false),
            }
        }
        Err(_) => assert!(true),
    }
}

#[test]
fn single_node() {
    let addr = get_storage_node(0, 1);
    let (insert_key, _) = key_and_value();

    {
        let content = InternodeRequest::Write {
            key: insert_key.to_owned(),
            value: Value::Value {
                content: vec![1],
                timestamp: 10000000,
            },
        };
        let addr = &addr.to_owned();
        let r: Future<InternodeResponse, Error> = client::Client::send_to_node(addr, &content);
        match r.await().unwrap() {
            InternodeResponse::WriteAck {key, timestamp} => {
                assert_eq!(key, insert_key);
                assert_eq!(timestamp, 10000000);
            },
            e => panic!("{:?}", e),
        }
    }
    {
        let content = InternodeRequest::Read {
            key: insert_key.to_owned(),
        };
        let addr = &addr.to_owned();
        let r: Future<InternodeResponse, Error> = client::Client::send_to_node(addr, &content);
        let r = r.await();
        match r {
            Ok(r) => match r {
                InternodeResponse::Value {key, value} => {
                    assert_eq!(key, insert_key);
                    match value {
                        Value::Value {content, timestamp} => {
                            assert_eq!(&content[..], &[1][..]);
                            assert_eq!(timestamp, 10000000);
                        },
                        _ => panic!(),
                    }
                },
                _ => panic!(),
            },
            Err(_) => panic!(),
        }
    }
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::net::{self, SocketAddr};
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
use std::sync::{Arc, Barrier};
use std::thread;

use mio::net::TcpListener;
use mio::{Interests, Poll, Token};

mod util;

use util::{
    any_local_address, any_local_ipv6_address, assert_send, assert_sync, assert_would_block,
    expect_events, expect_no_events, init, init_with_poll, ExpectEvent,
};

const ID1: Token = Token(0);
const ID2: Token = Token(1);

#[test]
fn is_send_and_sync() {
    assert_send::<TcpListener>();
    assert_sync::<TcpListener>();
}

#[test]
fn tcp_listener() {
    smoke_test_tcp_listener(any_local_address());
}

#[test]
fn tcp_listener_ipv6() {
    smoke_test_tcp_listener(any_local_ipv6_address());
}

fn smoke_test_tcp_listener(addr: SocketAddr) {
    let (mut poll, mut events) = init_with_poll();

    let listener = TcpListener::bind(addr).unwrap();
    let address = listener.local_addr().unwrap();

    poll.registry()
        .register(&listener, ID1, Interests::READABLE)
        .expect("unable to register TCP listener");

    let barrier = Arc::new(Barrier::new(2));
    let thread_handle = start_connections(address, 1, barrier.clone());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    // Expect a single connection.
    let (mut stream, peer_address) = listener.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    // Expect the stream to be non-blocking.
    let mut buf = [0; 20];
    assert_would_block(stream.read(&mut buf));

    // Expect no more connections.
    assert_would_block(listener.accept());

    assert!(listener.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn ttl() {
    init();

    let listener = TcpListener::bind(any_local_address()).unwrap();

    const TTL: u32 = 10;
    listener.set_ttl(TTL).unwrap();
    assert_eq!(listener.ttl().unwrap(), TTL);
    assert!(listener.take_error().unwrap().is_none());
}

#[test]
#[cfg(unix)]
fn raw_fd() {
    init();

    let listener = TcpListener::bind(any_local_address()).unwrap();
    let address = listener.local_addr().unwrap();

    let raw_fd1 = listener.as_raw_fd();
    let raw_fd2 = listener.into_raw_fd();
    assert_eq!(raw_fd1, raw_fd2);

    let listener = unsafe { TcpListener::from_raw_fd(raw_fd2) };
    assert_eq!(listener.as_raw_fd(), raw_fd1);
    assert_eq!(listener.local_addr().unwrap(), address);
}

#[test]
fn registering() {
    let (mut poll, mut events) = init_with_poll();

    let stream = TcpListener::bind(any_local_address()).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::READABLE)
        .expect("unable to register TCP listener");

    expect_no_events(&mut poll, &mut events);

    // NOTE: more tests are done in the smoke tests above.
}

#[test]
fn reregister() {
    let (mut poll, mut events) = init_with_poll();

    let listener = TcpListener::bind(any_local_address()).unwrap();
    let address = listener.local_addr().unwrap();

    poll.registry()
        .register(&listener, ID1, Interests::READABLE)
        .unwrap();
    poll.registry()
        .reregister(&listener, ID2, Interests::READABLE)
        .unwrap();

    let barrier = Arc::new(Barrier::new(2));
    let thread_handle = start_connections(address, 1, barrier.clone());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID2, Interests::READABLE)],
    );

    let (stream, peer_address) = listener.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    assert_would_block(listener.accept());

    assert!(listener.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle.join().expect("unable to join thread");
}

#[test]
#[cfg_attr(windows, ignore = "deregister doesn't work, see #1073")]
fn deregister() {
    let (mut poll, mut events) = init_with_poll();

    let listener = TcpListener::bind(any_local_address()).unwrap();
    let address = listener.local_addr().unwrap();

    poll.registry()
        .register(&listener, ID1, Interests::READABLE)
        .unwrap();

    let barrier = Arc::new(Barrier::new(2));
    let thread_handle = start_connections(address, 1, barrier.clone());

    poll.registry().deregister(&listener).unwrap();

    expect_no_events(&mut poll, &mut events);

    // Should still be able to accept the connection.
    let (stream, peer_address) = listener.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    assert_would_block(listener.accept());

    assert!(listener.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn try_clone_same_poll() {
    let (mut poll, mut events) = init_with_poll();

    let listener1 = TcpListener::bind(any_local_address()).unwrap();
    let listener2 = listener1.try_clone().expect("unable to clone TCP listener");
    #[cfg(unix)]
    assert_ne!(listener1.as_raw_fd(), listener2.as_raw_fd());
    let address = listener1.local_addr().unwrap();
    assert_eq!(address, listener2.local_addr().unwrap());

    let barrier = Arc::new(Barrier::new(3));
    let thread_handle1 = start_connections(address, 1, barrier.clone());
    let thread_handle2 = start_connections(address, 1, barrier.clone());

    poll.registry()
        .register(&listener1, ID1, Interests::READABLE)
        .unwrap();
    poll.registry()
        .register(&listener2, ID2, Interests::READABLE)
        .unwrap();

    expect_events(
        &mut poll,
        &mut events,
        vec![
            ExpectEvent::new(ID1, Interests::READABLE),
            ExpectEvent::new(ID2, Interests::READABLE),
        ],
    );

    let (stream, peer_address) = listener1.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    let (stream, peer_address) = listener2.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    assert_would_block(listener1.accept());
    assert_would_block(listener2.accept());

    assert!(listener1.take_error().unwrap().is_none());
    assert!(listener2.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle1.join().expect("unable to join thread");
    thread_handle2.join().expect("unable to join thread");
}

#[test]
fn try_clone_different_poll() {
    let (mut poll1, mut events) = init_with_poll();
    let mut poll2 = Poll::new().unwrap();

    let listener1 = TcpListener::bind(any_local_address()).unwrap();
    let listener2 = listener1.try_clone().expect("unable to clone TCP listener");
    #[cfg(unix)]
    assert_ne!(listener1.as_raw_fd(), listener2.as_raw_fd());
    let address = listener1.local_addr().unwrap();
    assert_eq!(address, listener2.local_addr().unwrap());

    let barrier = Arc::new(Barrier::new(3));
    let thread_handle1 = start_connections(address, 1, barrier.clone());
    let thread_handle2 = start_connections(address, 1, barrier.clone());

    poll1
        .registry()
        .register(&listener1, ID1, Interests::READABLE)
        .unwrap();
    poll2
        .registry()
        .register(&listener2, ID2, Interests::READABLE)
        .unwrap();

    expect_events(
        &mut poll1,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );
    expect_events(
        &mut poll2,
        &mut events,
        vec![ExpectEvent::new(ID2, Interests::READABLE)],
    );

    let (stream, peer_address) = listener1.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    let (stream, peer_address) = listener2.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    assert_would_block(listener1.accept());
    assert_would_block(listener2.accept());

    assert!(listener1.take_error().unwrap().is_none());
    assert!(listener2.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle1.join().expect("unable to join thread");
    thread_handle2.join().expect("unable to join thread");
}

/// Start `n_connections` connections to `address`. If a `barrier` is provided
/// it will wait on it after each connection is made before it is dropped.
fn start_connections(
    address: SocketAddr,
    n_connections: usize,
    barrier: Arc<Barrier>,
) -> thread::JoinHandle<()> {
    thread::spawn(move || {
        for _ in 0..n_connections {
            let conn = net::TcpStream::connect(address).unwrap();
            barrier.wait();
            drop(conn);
        }
    })
}
<commit_msg>Make TcpListener::try_clone tests more robust<commit_after>use std::io::Read;
use std::net::{self, SocketAddr};
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
use std::sync::{Arc, Barrier};
use std::thread;

use mio::net::TcpListener;
use mio::{Interests, Poll, Token};

mod util;

use util::{
    any_local_address, any_local_ipv6_address, assert_send, assert_sync, assert_would_block,
    expect_events, expect_no_events, init, init_with_poll, ExpectEvent,
};

const ID1: Token = Token(0);
const ID2: Token = Token(1);

#[test]
fn is_send_and_sync() {
    assert_send::<TcpListener>();
    assert_sync::<TcpListener>();
}

#[test]
fn tcp_listener() {
    smoke_test_tcp_listener(any_local_address());
}

#[test]
fn tcp_listener_ipv6() {
    smoke_test_tcp_listener(any_local_ipv6_address());
}

fn smoke_test_tcp_listener(addr: SocketAddr) {
    let (mut poll, mut events) = init_with_poll();

    let listener = TcpListener::bind(addr).unwrap();
    let address = listener.local_addr().unwrap();

    poll.registry()
        .register(&listener, ID1, Interests::READABLE)
        .expect("unable to register TCP listener");

    let barrier = Arc::new(Barrier::new(2));
    let thread_handle = start_connections(address, 1, barrier.clone());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    // Expect a single connection.
    let (mut stream, peer_address) = listener.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    // Expect the stream to be non-blocking.
    let mut buf = [0; 20];
    assert_would_block(stream.read(&mut buf));

    // Expect no more connections.
    assert_would_block(listener.accept());

    assert!(listener.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn ttl() {
    init();

    let listener = TcpListener::bind(any_local_address()).unwrap();

    const TTL: u32 = 10;
    listener.set_ttl(TTL).unwrap();
    assert_eq!(listener.ttl().unwrap(), TTL);
    assert!(listener.take_error().unwrap().is_none());
}

#[test]
#[cfg(unix)]
fn raw_fd() {
    init();

    let listener = TcpListener::bind(any_local_address()).unwrap();
    let address = listener.local_addr().unwrap();

    let raw_fd1 = listener.as_raw_fd();
    let raw_fd2 = listener.into_raw_fd();
    assert_eq!(raw_fd1, raw_fd2);

    let listener = unsafe { TcpListener::from_raw_fd(raw_fd2) };
    assert_eq!(listener.as_raw_fd(), raw_fd1);
    assert_eq!(listener.local_addr().unwrap(), address);
}

#[test]
fn registering() {
    let (mut poll, mut events) = init_with_poll();

    let stream = TcpListener::bind(any_local_address()).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::READABLE)
        .expect("unable to register TCP listener");

    expect_no_events(&mut poll, &mut events);

    // NOTE: more tests are done in the smoke tests above.
}

#[test]
fn reregister() {
    let (mut poll, mut events) = init_with_poll();

    let listener = TcpListener::bind(any_local_address()).unwrap();
    let address = listener.local_addr().unwrap();

    poll.registry()
        .register(&listener, ID1, Interests::READABLE)
        .unwrap();
    poll.registry()
        .reregister(&listener, ID2, Interests::READABLE)
        .unwrap();

    let barrier = Arc::new(Barrier::new(2));
    let thread_handle = start_connections(address, 1, barrier.clone());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID2, Interests::READABLE)],
    );

    let (stream, peer_address) = listener.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    assert_would_block(listener.accept());

    assert!(listener.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle.join().expect("unable to join thread");
}

#[test]
#[cfg_attr(windows, ignore = "deregister doesn't work, see #1073")]
fn deregister() {
    let (mut poll, mut events) = init_with_poll();

    let listener = TcpListener::bind(any_local_address()).unwrap();
    let address = listener.local_addr().unwrap();

    poll.registry()
        .register(&listener, ID1, Interests::READABLE)
        .unwrap();

    let barrier = Arc::new(Barrier::new(2));
    let thread_handle = start_connections(address, 1, barrier.clone());

    poll.registry().deregister(&listener).unwrap();

    expect_no_events(&mut poll, &mut events);

    // Should still be able to accept the connection.
    let (stream, peer_address) = listener.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    assert_would_block(listener.accept());

    assert!(listener.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn try_clone_same_poll() {
    let (mut poll, mut events) = init_with_poll();

    let listener1 = TcpListener::bind(any_local_address()).unwrap();
    let listener2 = listener1.try_clone().expect("unable to clone TCP listener");
    #[cfg(unix)]
    assert_ne!(listener1.as_raw_fd(), listener2.as_raw_fd());
    let address = listener1.local_addr().unwrap();
    assert_eq!(address, listener2.local_addr().unwrap());

    let barrier = Arc::new(Barrier::new(3));
    let thread_handle1 = start_connections(address, 1, barrier.clone());

    poll.registry()
        .register(&listener1, ID1, Interests::READABLE)
        .unwrap();
    poll.registry()
        .register(&listener2, ID2, Interests::READABLE)
        .unwrap();

    expect_events(
        &mut poll,
        &mut events,
        vec![
            ExpectEvent::new(ID1, Interests::READABLE),
            ExpectEvent::new(ID2, Interests::READABLE),
        ],
    );

    let (stream, peer_address) = listener1.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    let thread_handle2 = start_connections(address, 1, barrier.clone());

    expect_events(
        &mut poll,
        &mut events,
        vec![
            ExpectEvent::new(ID1, Interests::READABLE),
            ExpectEvent::new(ID2, Interests::READABLE),
        ],
    );

    let (stream, peer_address) = listener2.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    assert_would_block(listener1.accept());
    assert_would_block(listener2.accept());

    assert!(listener1.take_error().unwrap().is_none());
    assert!(listener2.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle1.join().expect("unable to join thread");
    thread_handle2.join().expect("unable to join thread");
}

#[test]
fn try_clone_different_poll() {
    let (mut poll1, mut events) = init_with_poll();
    let mut poll2 = Poll::new().unwrap();

    let listener1 = TcpListener::bind(any_local_address()).unwrap();
    let listener2 = listener1.try_clone().expect("unable to clone TCP listener");
    #[cfg(unix)]
    assert_ne!(listener1.as_raw_fd(), listener2.as_raw_fd());
    let address = listener1.local_addr().unwrap();
    assert_eq!(address, listener2.local_addr().unwrap());

    let barrier = Arc::new(Barrier::new(3));
    let thread_handle1 = start_connections(address, 1, barrier.clone());

    poll1
        .registry()
        .register(&listener1, ID1, Interests::READABLE)
        .unwrap();
    poll2
        .registry()
        .register(&listener2, ID2, Interests::READABLE)
        .unwrap();

    expect_events(
        &mut poll1,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );
    expect_events(
        &mut poll2,
        &mut events,
        vec![ExpectEvent::new(ID2, Interests::READABLE)],
    );

    let (stream, peer_address) = listener1.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    let thread_handle2 = start_connections(address, 1, barrier.clone());

    expect_events(
        &mut poll1,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );
    expect_events(
        &mut poll2,
        &mut events,
        vec![ExpectEvent::new(ID2, Interests::READABLE)],
    );

    let (stream, peer_address) = listener2.accept().expect("unable to accept connection");
    assert!(peer_address.ip().is_loopback());
    assert_eq!(stream.peer_addr().unwrap(), peer_address);
    assert_eq!(stream.local_addr().unwrap(), address);

    assert_would_block(listener1.accept());
    assert_would_block(listener2.accept());

    assert!(listener1.take_error().unwrap().is_none());
    assert!(listener2.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle1.join().expect("unable to join thread");
    thread_handle2.join().expect("unable to join thread");
}

/// Start `n_connections` connections to `address`. If a `barrier` is provided
/// it will wait on it after each connection is made before it is dropped.
fn start_connections(
    address: SocketAddr,
    n_connections: usize,
    barrier: Arc<Barrier>,
) -> thread::JoinHandle<()> {
    thread::spawn(move || {
        for _ in 0..n_connections {
            let conn = net::TcpStream::connect(address).unwrap();
            barrier.wait();
            drop(conn);
        }
    })
}
<|endoftext|>


<commit_before>use crate::config::{Config, Solver};
use crate::heuristic;
use oppai_field::field::{self, Field, NonZeroPos};
use oppai_field::player::Player;
use oppai_field::zobrist::Zobrist;
use oppai_ladders::ladders::ladders;
use oppai_minimax::minimax::Minimax;
use oppai_patterns::patterns::Patterns;
use oppai_uct::uct::UctRoot;
use rand::distributions::{Distribution, Standard};
use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
use std::{
  cmp,
  sync::atomic::{AtomicBool, Ordering},
  sync::Arc,
  time::{Duration, Instant},
};

fn is_field_occupied(field: &Field) -> bool {
  for pos in field.min_pos()..=field.max_pos() {
    if field.cell(pos).is_putting_allowed() {
      return false;
    }
  }
  true
}

fn with_timeout<T: Send, F: FnOnce() -> T + Send>(f: F, should_stop: &AtomicBool, time: Duration) -> T {
  let (s, r) = crossbeam::channel::bounded(1);
  crossbeam::scope(|scope| {
    scope.spawn(move |_| {
      let result = f();
      s.send(result).unwrap();
    });
    if let Ok(result) = r.recv_timeout(time) {
      return result;
    }
    should_stop.store(true, Ordering::Relaxed);
    r.recv().unwrap()
  })
  .unwrap()
}

pub struct Bot<R> {
  pub rng: R,
  pub patterns: Arc<Patterns>,
  pub field: Field,
  pub uct: UctRoot,
  pub minimax: Minimax,
  pub config: Config,
}

impl<S, R> Bot<R>
where
  R: Rng + SeedableRng<Seed = S> + Send,
  Standard: Distribution<S>,
{
  pub fn new(width: u32, height: u32, mut rng: R, patterns: Arc<Patterns>, config: Config) -> Self {
    info!("Initialization with width {0}, height {1}.", width, height);
    let length = field::length(width, height);
    let zobrist = Arc::new(Zobrist::new(length * 2, &mut rng));
    let field_zobrist = Arc::clone(&zobrist);
    Bot {
      rng,
      patterns,
      field: Field::new(width, height, field_zobrist),
      uct: UctRoot::new(config.uct.clone(), length),
      minimax: Minimax::new(config.minimax.clone()),
      config,
    }
  }

  pub fn initial_move(&self) -> Option<NonZeroPos> {
    let result = match self.field.moves_count() {
      0 => Some((self.field.width() / 2, self.field.height() / 2)),
      1 => {
        let width = self.field.width();
        let height = self.field.height();
        let pos = self.field.points_seq()[0];
        let x = self.field.to_x(pos);
        let y = self.field.to_y(pos);
        if x == 0 || x == width - 1 || y == 0 || y == height - 1 {
          Some((width / 2, height / 2))
        } else if cmp::min(x, width - x - 1) < cmp::min(y, height - y - 1) {
          if x < width - x - 1 {
            Some((x + 1, y))
          } else {
            Some((x - 1, y))
          }
        } else if cmp::min(x, width - x - 1) > cmp::min(y, height - y - 1) {
          if y < height - y - 1 {
            Some((x, y + 1))
          } else {
            Some((x, y - 1))
          }
        } else {
          let dx = x as i32 - (width / 2) as i32;
          let dy = y as i32 - (height / 2) as i32;
          if dx.abs() > dy.abs() {
            if dx < 0 {
              Some((x + 1, y))
            } else {
              Some((x - 1, y))
            }
          } else if dy < 0 {
            Some((x, y + 1))
          } else {
            Some((x, y - 1))
          }
        }
      }
      _ => None,
    };
    result.and_then(|(x, y)| NonZeroPos::new(self.field.to_pos(x, y)))
  }

  pub fn best_move(
    &mut self,
    player: Player,
    uct_iterations: usize,
    minimax_depth: u32,
    should_stop: &AtomicBool,
  ) -> Option<NonZeroPos> {
    let now = Instant::now();
    if self.field.width() < 3 || self.field.height() < 3 || is_field_occupied(&self.field) {
      return None;
    }
    if let Some(pos) = self.initial_move() {
      return Some(pos);
    }
    if let Some(&pos) = self.patterns.find(&self.field, player, false).choose(&mut self.rng) {
      info!(
        "Cumulative time for patterns evaluation (move is found): {:?}.",
        now.elapsed()
      );
      return NonZeroPos::new(pos);
    }
    let elapsed = now.elapsed();
    info!("Cumulative time for patterns evaluation: {:?}.", elapsed);
    if self.config.ladders {
      let ladders_time_limit = self.config.ladders_time_limit;
      if let (Some(pos), score, depth) = with_timeout(
        || ladders(&mut self.field, player, &should_stop),
        &should_stop,
        ladders_time_limit,
      ) {
        info!(
          "Cumulative time for ladders evaluation (move is found with score {} and depth {}): {:?}.",
          score,
          depth,
          now.elapsed()
        );
        if (score - self.field.score(player)) as u32 > self.config.ladders_score_limit
          && depth > self.config.ladders_depth_limit
        {
          return Some(pos);
        }
      };
    }
    let elapsed = now.elapsed();
    info!("Cumulative time for ladders evaluation: {:?}.", elapsed);
    let result = match self.config.solver {
      Solver::Uct => self
        .uct
        .best_move(&self.field, player, &mut self.rng, &should_stop, uct_iterations)
        .or_else(|| heuristic::heuristic(&self.field, player)),
      Solver::Minimax => self
        .minimax
        .minimax(&mut self.field, player, minimax_depth, &should_stop)
        .or_else(|| heuristic::heuristic(&self.field, player)),
      Solver::Heuristic => heuristic::heuristic(&self.field, player),
    };
    info!("Cumulative time for best move evaluation: {:?}", now.elapsed());
    result
  }

  pub fn best_move_with_time(
    &mut self,
    player: Player,
    time: Duration,
    should_stop: &AtomicBool,
  ) -> Option<NonZeroPos> {
    let now = Instant::now();
    let time = time - self.config.time_gap;
    if self.field.width() < 3 || self.field.height() < 3 || is_field_occupied(&self.field) {
      return None;
    }
    if let Some(pos) = self.initial_move() {
      return Some(pos);
    }
    if let Some(&pos) = self.patterns.find(&self.field, player, false).choose(&mut self.rng) {
      info!(
        "Cumulative time for patterns evaluation (move is found): {:?}.",
        now.elapsed()
      );
      return NonZeroPos::new(pos);
    }
    let elapsed = now.elapsed();
    info!("Cumulative time for patterns evaluation: {:?}.", elapsed);
    let time_left = if let Some(time_left) = time.checked_sub(elapsed) {
      time_left
    } else {
      return None;
    };
    if self.config.ladders {
      let ladders_time_limit = self.config.ladders_time_limit;
      if let (Some(pos), score, depth) = with_timeout(
        || ladders(&mut self.field, player, &should_stop),
        &should_stop,
        ladders_time_limit.min(time_left),
      ) {
        info!(
          "Cumulative time for ladders evaluation (move is found with score {} and depth {}): {:?}.",
          score,
          depth,
          now.elapsed()
        );
        if (score - self.field.score(player)) as u32 > self.config.ladders_score_limit
          && depth > self.config.ladders_depth_limit
        {
          return Some(pos);
        }
      };
    }
    let elapsed = now.elapsed();
    info!("Cumulative time for ladders evaluation: {:?}.", elapsed);
    let time_left = if let Some(time_left) = time.checked_sub(elapsed) {
      time_left
    } else {
      return None;
    };
    let result = match self.config.solver {
      Solver::Uct => with_timeout(
        || {
          self
            .uct
            .best_move(&self.field, player, &mut self.rng, &should_stop, usize::max_value())
            .or_else(|| heuristic::heuristic(&self.field, player))
        },
        &should_stop,
        time_left,
      ),
      Solver::Minimax => with_timeout(
        || {
          self
            .minimax
            .minimax_with_time(&mut self.field, player, &should_stop)
            .or_else(|| heuristic::heuristic(&self.field, player))
        },
        &should_stop,
        time_left,
      ),
      Solver::Heuristic => heuristic::heuristic(&self.field, player),
    };
    info!("Cumulative time for best move evaluation: {:?}", now.elapsed());
    result
  }
}
<commit_msg>Use heuristic in case of timeout.<commit_after>use crate::config::{Config, Solver};
use crate::heuristic;
use oppai_field::field::{self, Field, NonZeroPos};
use oppai_field::player::Player;
use oppai_field::zobrist::Zobrist;
use oppai_ladders::ladders::ladders;
use oppai_minimax::minimax::Minimax;
use oppai_patterns::patterns::Patterns;
use oppai_uct::uct::UctRoot;
use rand::distributions::{Distribution, Standard};
use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};
use std::{
  cmp,
  sync::atomic::{AtomicBool, Ordering},
  sync::Arc,
  time::{Duration, Instant},
};

fn is_field_occupied(field: &Field) -> bool {
  for pos in field.min_pos()..=field.max_pos() {
    if field.cell(pos).is_putting_allowed() {
      return false;
    }
  }
  true
}

fn with_timeout<T: Send, F: FnOnce() -> T + Send>(f: F, should_stop: &AtomicBool, time: Duration) -> T {
  let (s, r) = crossbeam::channel::bounded(1);
  crossbeam::scope(|scope| {
    scope.spawn(move |_| {
      let result = f();
      s.send(result).unwrap();
    });
    if let Ok(result) = r.recv_timeout(time) {
      return result;
    }
    should_stop.store(true, Ordering::Relaxed);
    r.recv().unwrap()
  })
  .unwrap()
}

pub struct Bot<R> {
  pub rng: R,
  pub patterns: Arc<Patterns>,
  pub field: Field,
  pub uct: UctRoot,
  pub minimax: Minimax,
  pub config: Config,
}

impl<S, R> Bot<R>
where
  R: Rng + SeedableRng<Seed = S> + Send,
  Standard: Distribution<S>,
{
  pub fn new(width: u32, height: u32, mut rng: R, patterns: Arc<Patterns>, config: Config) -> Self {
    info!("Initialization with width {0}, height {1}.", width, height);
    let length = field::length(width, height);
    let zobrist = Arc::new(Zobrist::new(length * 2, &mut rng));
    let field_zobrist = Arc::clone(&zobrist);
    Bot {
      rng,
      patterns,
      field: Field::new(width, height, field_zobrist),
      uct: UctRoot::new(config.uct.clone(), length),
      minimax: Minimax::new(config.minimax.clone()),
      config,
    }
  }

  pub fn initial_move(&self) -> Option<NonZeroPos> {
    let result = match self.field.moves_count() {
      0 => Some((self.field.width() / 2, self.field.height() / 2)),
      1 => {
        let width = self.field.width();
        let height = self.field.height();
        let pos = self.field.points_seq()[0];
        let x = self.field.to_x(pos);
        let y = self.field.to_y(pos);
        if x == 0 || x == width - 1 || y == 0 || y == height - 1 {
          Some((width / 2, height / 2))
        } else if cmp::min(x, width - x - 1) < cmp::min(y, height - y - 1) {
          if x < width - x - 1 {
            Some((x + 1, y))
          } else {
            Some((x - 1, y))
          }
        } else if cmp::min(x, width - x - 1) > cmp::min(y, height - y - 1) {
          if y < height - y - 1 {
            Some((x, y + 1))
          } else {
            Some((x, y - 1))
          }
        } else {
          let dx = x as i32 - (width / 2) as i32;
          let dy = y as i32 - (height / 2) as i32;
          if dx.abs() > dy.abs() {
            if dx < 0 {
              Some((x + 1, y))
            } else {
              Some((x - 1, y))
            }
          } else if dy < 0 {
            Some((x, y + 1))
          } else {
            Some((x, y - 1))
          }
        }
      }
      _ => None,
    };
    result.and_then(|(x, y)| NonZeroPos::new(self.field.to_pos(x, y)))
  }

  pub fn best_move(
    &mut self,
    player: Player,
    uct_iterations: usize,
    minimax_depth: u32,
    should_stop: &AtomicBool,
  ) -> Option<NonZeroPos> {
    let now = Instant::now();
    if self.field.width() < 3 || self.field.height() < 3 || is_field_occupied(&self.field) {
      return None;
    }
    if let Some(pos) = self.initial_move() {
      return Some(pos);
    }
    if let Some(&pos) = self.patterns.find(&self.field, player, false).choose(&mut self.rng) {
      info!(
        "Cumulative time for patterns evaluation (move is found): {:?}.",
        now.elapsed()
      );
      return NonZeroPos::new(pos);
    }
    let elapsed = now.elapsed();
    info!("Cumulative time for patterns evaluation: {:?}.", elapsed);
    if self.config.ladders {
      let ladders_time_limit = self.config.ladders_time_limit;
      if let (Some(pos), score, depth) = with_timeout(
        || ladders(&mut self.field, player, &should_stop),
        &should_stop,
        ladders_time_limit,
      ) {
        info!(
          "Cumulative time for ladders evaluation (move is found with score {} and depth {}): {:?}.",
          score,
          depth,
          now.elapsed()
        );
        if (score - self.field.score(player)) as u32 > self.config.ladders_score_limit
          && depth > self.config.ladders_depth_limit
        {
          return Some(pos);
        }
      };
    }
    let elapsed = now.elapsed();
    info!("Cumulative time for ladders evaluation: {:?}.", elapsed);
    let result = match self.config.solver {
      Solver::Uct => self
        .uct
        .best_move(&self.field, player, &mut self.rng, &should_stop, uct_iterations)
        .or_else(|| heuristic::heuristic(&self.field, player)),
      Solver::Minimax => self
        .minimax
        .minimax(&mut self.field, player, minimax_depth, &should_stop)
        .or_else(|| heuristic::heuristic(&self.field, player)),
      Solver::Heuristic => heuristic::heuristic(&self.field, player),
    };
    info!("Cumulative time for best move evaluation: {:?}", now.elapsed());
    result
  }

  pub fn best_move_with_time(
    &mut self,
    player: Player,
    time: Duration,
    should_stop: &AtomicBool,
  ) -> Option<NonZeroPos> {
    let now = Instant::now();
    let time = time - self.config.time_gap;
    if self.field.width() < 3 || self.field.height() < 3 || is_field_occupied(&self.field) {
      return None;
    }
    if let Some(pos) = self.initial_move() {
      return Some(pos);
    }
    if let Some(&pos) = self.patterns.find(&self.field, player, false).choose(&mut self.rng) {
      info!(
        "Cumulative time for patterns evaluation (move is found): {:?}.",
        now.elapsed()
      );
      return NonZeroPos::new(pos);
    }
    let elapsed = now.elapsed();
    info!("Cumulative time for patterns evaluation: {:?}.", elapsed);
    let time_left = if let Some(time_left) = time.checked_sub(elapsed) {
      time_left
    } else {
      return heuristic::heuristic(&self.field, player);
    };
    if self.config.ladders {
      let ladders_time_limit = self.config.ladders_time_limit;
      if let (Some(pos), score, depth) = with_timeout(
        || ladders(&mut self.field, player, &should_stop),
        &should_stop,
        ladders_time_limit.min(time_left),
      ) {
        info!(
          "Cumulative time for ladders evaluation (move is found with score {} and depth {}): {:?}.",
          score,
          depth,
          now.elapsed()
        );
        if (score - self.field.score(player)) as u32 > self.config.ladders_score_limit
          && depth > self.config.ladders_depth_limit
        {
          return Some(pos);
        }
      };
    }
    let elapsed = now.elapsed();
    info!("Cumulative time for ladders evaluation: {:?}.", elapsed);
    let time_left = if let Some(time_left) = time.checked_sub(elapsed) {
      time_left
    } else {
      return heuristic::heuristic(&self.field, player);
    };
    let result = match self.config.solver {
      Solver::Uct => with_timeout(
        || {
          self
            .uct
            .best_move(&self.field, player, &mut self.rng, &should_stop, usize::max_value())
            .or_else(|| heuristic::heuristic(&self.field, player))
        },
        &should_stop,
        time_left,
      ),
      Solver::Minimax => with_timeout(
        || {
          self
            .minimax
            .minimax_with_time(&mut self.field, player, &should_stop)
            .or_else(|| heuristic::heuristic(&self.field, player))
        },
        &should_stop,
        time_left,
      ),
      Solver::Heuristic => heuristic::heuristic(&self.field, player),
    };
    info!("Cumulative time for best move evaluation: {:?}", now.elapsed());
    result
  }
}
<|endoftext|>


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

use proc_macro2::Span;
use syn::Result as ParseResult;
use syn::{
    parse::{Parse, ParseStream},
    Token,
};

use crate::type_config::TypeConfig;

#[derive(PartialEq, Clone, Debug, Hash)]
pub enum UnsafePolicy {
    AllFunctionsSafe,
    AllFunctionsUnsafe,
}

impl Parse for UnsafePolicy {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        if input.parse::<Option<Token![unsafe]>>()?.is_some() {
            return Ok(UnsafePolicy::AllFunctionsSafe);
        }
        let r = match input.parse::<Option<syn::Ident>>()? {
            Some(id) => {
                if id == "unsafe_ffi" {
                    Ok(UnsafePolicy::AllFunctionsSafe)
                } else {
                    Err(syn::Error::new(id.span(), "expected unsafe_ffi"))
                }
            }
            None => Ok(UnsafePolicy::AllFunctionsUnsafe),
        };
        if !input.is_empty() {
            return Err(syn::Error::new(
                Span::call_site(),
                "unexpected tokens within safety directive",
            ));
        }
        r
    }
}

#[derive(Hash, Debug)]
pub struct IncludeCppConfig {
    pub inclusions: Vec<String>,
    pub unsafe_policy: UnsafePolicy,
    pub type_config: TypeConfig,
    pub parse_only: bool,
}

impl Parse for IncludeCppConfig {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        // Takes as inputs:
        // 1. List of headers to include
        // 2. List of #defines to include
        // 3. Allowlist

        let mut inclusions = Vec::new();
        let mut parse_only = false;
        let mut type_config = TypeConfig::new();
        let mut unsafe_policy = UnsafePolicy::AllFunctionsUnsafe;

        while !input.is_empty() {
            let has_hexathorpe = input.parse::<Option<syn::Token![#]>>()?.is_some();
            let ident: syn::Ident = input.parse()?;
            if has_hexathorpe {
                if ident != "include" {
                    return Err(syn::Error::new(ident.span(), "expected include"));
                }
                let hdr: syn::LitStr = input.parse()?;
                inclusions.push(hdr.value());
            } else {
                input.parse::<Option<syn::Token![!]>>()?;
                if ident == "generate" || ident == "generate_pod" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    type_config.allowlist.push(generate.value());
                    if ident == "generate_pod" {
                        type_config.pod_requests.push(generate.value());
                    }
                } else if ident == "block" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    type_config.blocklist.push(generate.value());
                } else if ident == "parse_only" {
                    parse_only = true;
                } else if ident == "exclude_utilities" {
                    type_config.exclude_utilities = true;
                } else if ident == "safety" {
                    let args;
                    syn::parenthesized!(args in input);
                    unsafe_policy = args.parse()?;
                } else {
                    return Err(syn::Error::new(
                        ident.span(),
                        "expected generate, generate_pod, nested_type, safety or exclude_utilities",
                    ));
                }
            }
            if input.is_empty() {
                break;
            }
        }

        Ok(IncludeCppConfig {
            inclusions,
            unsafe_policy,
            type_config,
            parse_only,
        })
    }
}

#[cfg(test)]
mod parse_tests {
    use crate::config::UnsafePolicy;
    use syn::parse_quote;
    #[test]
    fn test_safety_unsafe() {
        let us: UnsafePolicy = parse_quote! {
            unsafe
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_unsafe_ffi() {
        let us: UnsafePolicy = parse_quote! {
            unsafe_ffi
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_safe() {
        let us: UnsafePolicy = parse_quote! {};
        assert_eq!(us, UnsafePolicy::AllFunctionsUnsafe)
    }
}
<commit_msg>Swallowing parethenses after nullary directives.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use proc_macro2::Span;
use syn::{
    parse::{Parse, ParseStream},
    Token,
};
use syn::{Ident, Result as ParseResult};

use crate::type_config::TypeConfig;

#[derive(PartialEq, Clone, Debug, Hash)]
pub enum UnsafePolicy {
    AllFunctionsSafe,
    AllFunctionsUnsafe,
}

impl Parse for UnsafePolicy {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        if input.parse::<Option<Token![unsafe]>>()?.is_some() {
            return Ok(UnsafePolicy::AllFunctionsSafe);
        }
        let r = match input.parse::<Option<syn::Ident>>()? {
            Some(id) => {
                if id == "unsafe_ffi" {
                    Ok(UnsafePolicy::AllFunctionsSafe)
                } else {
                    Err(syn::Error::new(id.span(), "expected unsafe_ffi"))
                }
            }
            None => Ok(UnsafePolicy::AllFunctionsUnsafe),
        };
        if !input.is_empty() {
            return Err(syn::Error::new(
                Span::call_site(),
                "unexpected tokens within safety directive",
            ));
        }
        r
    }
}

#[derive(Hash, Debug)]
pub struct IncludeCppConfig {
    pub inclusions: Vec<String>,
    pub unsafe_policy: UnsafePolicy,
    pub type_config: TypeConfig,
    pub parse_only: bool,
}

impl Parse for IncludeCppConfig {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        // Takes as inputs:
        // 1. List of headers to include
        // 2. List of #defines to include
        // 3. Allowlist

        let mut inclusions = Vec::new();
        let mut parse_only = false;
        let mut type_config = TypeConfig::new();
        let mut unsafe_policy = UnsafePolicy::AllFunctionsUnsafe;

        while !input.is_empty() {
            let has_hexathorpe = input.parse::<Option<syn::Token![#]>>()?.is_some();
            let ident: syn::Ident = input.parse()?;
            if has_hexathorpe {
                if ident != "include" {
                    return Err(syn::Error::new(ident.span(), "expected include"));
                }
                let hdr: syn::LitStr = input.parse()?;
                inclusions.push(hdr.value());
            } else {
                input.parse::<Option<syn::Token![!]>>()?;
                if ident == "generate" || ident == "generate_pod" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    type_config.allowlist.push(generate.value());
                    if ident == "generate_pod" {
                        type_config.pod_requests.push(generate.value());
                    }
                } else if ident == "block" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    type_config.blocklist.push(generate.value());
                } else if ident == "parse_only" {
                    parse_only = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "exclude_utilities" {
                    type_config.exclude_utilities = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "safety" {
                    let args;
                    syn::parenthesized!(args in input);
                    unsafe_policy = args.parse()?;
                } else {
                    return Err(syn::Error::new(
                        ident.span(),
                        "expected generate, generate_pod, nested_type, safety or exclude_utilities",
                    ));
                }
            }
            if input.is_empty() {
                break;
            }
        }

        Ok(IncludeCppConfig {
            inclusions,
            unsafe_policy,
            type_config,
            parse_only,
        })
    }
}

fn swallow_parentheses(input: &ParseStream, latest_ident: &Ident) -> ParseResult<()> {
    let args;
    syn::parenthesized!(args in input);
    if args.is_empty() {
        Ok(())
    } else {
        Err(syn::Error::new(
            latest_ident.span(),
            "expected no arguments to directive",
        ))
    }
}

#[cfg(test)]
mod parse_tests {
    use crate::config::UnsafePolicy;
    use syn::parse_quote;
    #[test]
    fn test_safety_unsafe() {
        let us: UnsafePolicy = parse_quote! {
            unsafe
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_unsafe_ffi() {
        let us: UnsafePolicy = parse_quote! {
            unsafe_ffi
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_safe() {
        let us: UnsafePolicy = parse_quote! {};
        assert_eq!(us, UnsafePolicy::AllFunctionsUnsafe)
    }
}
<|endoftext|>


<commit_before>use crate::future;

use std::io;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

/// Converts or resolves without blocking to one or more `SocketAddr` values.
///
/// # DNS
///
/// Implementations of `ToSocketAddrs` for string types require a DNS lookup.
/// These implementations are only provided when Tokio is used with the
/// **`dns`** feature flag.
///
/// # Calling
///
/// Currently, this trait is only used as an argument to Tokio functions that
/// need to reference a target socket address. To perform a `SocketAddr`
/// conversion directly, use [`lookup_host()`](super::lookup_host()).
///
/// This trait is sealed and is intended to be opaque. The details of the trait
/// will change. Stabilization is pending enhancements to the Rust langague.
pub trait ToSocketAddrs: sealed::ToSocketAddrsPriv {}

type ReadyFuture<T> = future::Ready<io::Result<T>>;

// ===== impl &impl ToSocketAddrs =====

impl<T: ToSocketAddrs + ?Sized> ToSocketAddrs for &T {}

impl<T> sealed::ToSocketAddrsPriv for &T
where
    T: sealed::ToSocketAddrsPriv + ?Sized,
{
    type Iter = T::Iter;
    type Future = T::Future;

    fn to_socket_addrs(&self) -> Self::Future {
        (**self).to_socket_addrs()
    }
}

// ===== impl SocketAddr =====

impl ToSocketAddrs for SocketAddr {}

impl sealed::ToSocketAddrsPriv for SocketAddr {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        let iter = Some(*self).into_iter();
        future::ok(iter)
    }
}

// ===== impl SocketAddrV4 =====

impl ToSocketAddrs for SocketAddrV4 {}

impl sealed::ToSocketAddrsPriv for SocketAddrV4 {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        SocketAddr::V4(*self).to_socket_addrs()
    }
}

// ===== impl SocketAddrV6 =====

impl ToSocketAddrs for SocketAddrV6 {}

impl sealed::ToSocketAddrsPriv for SocketAddrV6 {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        SocketAddr::V6(*self).to_socket_addrs()
    }
}

// ===== impl (IpAddr, u16) =====

impl ToSocketAddrs for (IpAddr, u16) {}

impl sealed::ToSocketAddrsPriv for (IpAddr, u16) {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        let iter = Some(SocketAddr::from(*self)).into_iter();
        future::ok(iter)
    }
}

// ===== impl (Ipv4Addr, u16) =====

impl ToSocketAddrs for (Ipv4Addr, u16) {}

impl sealed::ToSocketAddrsPriv for (Ipv4Addr, u16) {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        let (ip, port) = *self;
        SocketAddrV4::new(ip, port).to_socket_addrs()
    }
}

// ===== impl (Ipv6Addr, u16) =====

impl ToSocketAddrs for (Ipv6Addr, u16) {}

impl sealed::ToSocketAddrsPriv for (Ipv6Addr, u16) {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        let (ip, port) = *self;
        SocketAddrV6::new(ip, port, 0, 0).to_socket_addrs()
    }
}

cfg_dns! {
    // ===== impl str =====

    impl ToSocketAddrs for str {}

    impl sealed::ToSocketAddrsPriv for str {
        type Iter = sealed::OneOrMore;
        type Future = sealed::MaybeReady;

        fn to_socket_addrs(&self) -> Self::Future {
            use crate::runtime::spawn_blocking;
            use sealed::MaybeReady;

            // First check if the input parses as a socket address
            let res: Result<SocketAddr, _> = self.parse();

            if let Ok(addr) = res {
                return MaybeReady::Ready(Some(addr));
            }

            // Run DNS lookup on the blocking pool
            let s = self.to_owned();

            MaybeReady::Blocking(spawn_blocking(move || {
                std::net::ToSocketAddrs::to_socket_addrs(&s)
            }))
        }
    }

    // ===== impl (&str, u16) =====

    impl ToSocketAddrs for (&str, u16) {}

    impl sealed::ToSocketAddrsPriv for (&str, u16) {
        type Iter = sealed::OneOrMore;
        type Future = sealed::MaybeReady;

        fn to_socket_addrs(&self) -> Self::Future {
            use crate::runtime::spawn_blocking;
            use sealed::MaybeReady;

            let (host, port) = *self;

            // try to parse the host as a regular IP address first
            if let Ok(addr) = host.parse::<Ipv4Addr>() {
                let addr = SocketAddrV4::new(addr, port);
                let addr = SocketAddr::V4(addr);

                return MaybeReady::Ready(Some(addr));
            }

            if let Ok(addr) = host.parse::<Ipv6Addr>() {
                let addr = SocketAddrV6::new(addr, port, 0, 0);
                let addr = SocketAddr::V6(addr);

                return MaybeReady::Ready(Some(addr));
            }

            let host = host.to_owned();

            MaybeReady::Blocking(spawn_blocking(move || {
                std::net::ToSocketAddrs::to_socket_addrs(&(&host[..], port))
            }))
        }
    }

    // ===== impl String =====

    impl ToSocketAddrs for String {}

    impl sealed::ToSocketAddrsPriv for String {
        type Iter = <str as sealed::ToSocketAddrsPriv>::Iter;
        type Future = <str as sealed::ToSocketAddrsPriv>::Future;

        fn to_socket_addrs(&self) -> Self::Future {
            (&self[..]).to_socket_addrs()
        }
    }
}

pub(crate) mod sealed {
    //! The contents of this trait are intended to remain private and __not__
    //! part of the `ToSocketAddrs` public API. The details will change over
    //! time.

    use std::future::Future;
    use std::io;
    use std::net::SocketAddr;

    cfg_dns! {
        use crate::task::JoinHandle;

        use std::option;
        use std::pin::Pin;
        use std::task::{Context, Poll};
        use std::vec;
    }

    #[doc(hidden)]
    pub trait ToSocketAddrsPriv {
        type Iter: Iterator<Item = SocketAddr> + Send + 'static;
        type Future: Future<Output = io::Result<Self::Iter>> + Send + 'static;

        fn to_socket_addrs(&self) -> Self::Future;
    }

    cfg_dns! {
        #[doc(hidden)]
        #[derive(Debug)]
        pub enum MaybeReady {
            Ready(Option<SocketAddr>),
            Blocking(JoinHandle<io::Result<vec::IntoIter<SocketAddr>>>),
        }

        #[doc(hidden)]
        #[derive(Debug)]
        pub enum OneOrMore {
            One(option::IntoIter<SocketAddr>),
            More(vec::IntoIter<SocketAddr>),
        }

        impl Future for MaybeReady {
            type Output = io::Result<OneOrMore>;

            fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                match *self {
                    MaybeReady::Ready(ref mut i) => {
                        let iter = OneOrMore::One(i.take().into_iter());
                        Poll::Ready(Ok(iter))
                    }
                    MaybeReady::Blocking(ref mut rx) => {
                        let res = ready!(Pin::new(rx).poll(cx))?.map(OneOrMore::More);

                        Poll::Ready(res)
                    }
                }
            }
        }

        impl Iterator for OneOrMore {
            type Item = SocketAddr;

            fn next(&mut self) -> Option<Self::Item> {
                match self {
                    OneOrMore::One(i) => i.next(),
                    OneOrMore::More(i) => i.next(),
                }
            }

            fn size_hint(&self) -> (usize, Option<usize>) {
                match self {
                    OneOrMore::One(i) => i.size_hint(),
                    OneOrMore::More(i) => i.size_hint(),
                }
            }
        }
    }
}
<commit_msg>net: impl ToSocketAddrs for &[SocketAddr] (#2604)<commit_after>use crate::future;

use std::io;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

/// Converts or resolves without blocking to one or more `SocketAddr` values.
///
/// # DNS
///
/// Implementations of `ToSocketAddrs` for string types require a DNS lookup.
/// These implementations are only provided when Tokio is used with the
/// **`dns`** feature flag.
///
/// # Calling
///
/// Currently, this trait is only used as an argument to Tokio functions that
/// need to reference a target socket address. To perform a `SocketAddr`
/// conversion directly, use [`lookup_host()`](super::lookup_host()).
///
/// This trait is sealed and is intended to be opaque. The details of the trait
/// will change. Stabilization is pending enhancements to the Rust langague.
pub trait ToSocketAddrs: sealed::ToSocketAddrsPriv {}

type ReadyFuture<T> = future::Ready<io::Result<T>>;

// ===== impl &impl ToSocketAddrs =====

impl<T: ToSocketAddrs + ?Sized> ToSocketAddrs for &T {}

impl<T> sealed::ToSocketAddrsPriv for &T
where
    T: sealed::ToSocketAddrsPriv + ?Sized,
{
    type Iter = T::Iter;
    type Future = T::Future;

    fn to_socket_addrs(&self) -> Self::Future {
        (**self).to_socket_addrs()
    }
}

// ===== impl SocketAddr =====

impl ToSocketAddrs for SocketAddr {}

impl sealed::ToSocketAddrsPriv for SocketAddr {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        let iter = Some(*self).into_iter();
        future::ok(iter)
    }
}

// ===== impl SocketAddrV4 =====

impl ToSocketAddrs for SocketAddrV4 {}

impl sealed::ToSocketAddrsPriv for SocketAddrV4 {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        SocketAddr::V4(*self).to_socket_addrs()
    }
}

// ===== impl SocketAddrV6 =====

impl ToSocketAddrs for SocketAddrV6 {}

impl sealed::ToSocketAddrsPriv for SocketAddrV6 {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        SocketAddr::V6(*self).to_socket_addrs()
    }
}

// ===== impl (IpAddr, u16) =====

impl ToSocketAddrs for (IpAddr, u16) {}

impl sealed::ToSocketAddrsPriv for (IpAddr, u16) {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        let iter = Some(SocketAddr::from(*self)).into_iter();
        future::ok(iter)
    }
}

// ===== impl (Ipv4Addr, u16) =====

impl ToSocketAddrs for (Ipv4Addr, u16) {}

impl sealed::ToSocketAddrsPriv for (Ipv4Addr, u16) {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        let (ip, port) = *self;
        SocketAddrV4::new(ip, port).to_socket_addrs()
    }
}

// ===== impl (Ipv6Addr, u16) =====

impl ToSocketAddrs for (Ipv6Addr, u16) {}

impl sealed::ToSocketAddrsPriv for (Ipv6Addr, u16) {
    type Iter = std::option::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        let (ip, port) = *self;
        SocketAddrV6::new(ip, port, 0, 0).to_socket_addrs()
    }
}

// ===== impl &[SocketAddr] =====

impl ToSocketAddrs for &[SocketAddr] {}

impl sealed::ToSocketAddrsPriv for &[SocketAddr] {
    type Iter = std::vec::IntoIter<SocketAddr>;
    type Future = ReadyFuture<Self::Iter>;

    fn to_socket_addrs(&self) -> Self::Future {
        let iter = self.to_vec().into_iter();
        future::ok(iter)
    }
}

cfg_dns! {
    // ===== impl str =====

    impl ToSocketAddrs for str {}

    impl sealed::ToSocketAddrsPriv for str {
        type Iter = sealed::OneOrMore;
        type Future = sealed::MaybeReady;

        fn to_socket_addrs(&self) -> Self::Future {
            use crate::runtime::spawn_blocking;
            use sealed::MaybeReady;

            // First check if the input parses as a socket address
            let res: Result<SocketAddr, _> = self.parse();

            if let Ok(addr) = res {
                return MaybeReady::Ready(Some(addr));
            }

            // Run DNS lookup on the blocking pool
            let s = self.to_owned();

            MaybeReady::Blocking(spawn_blocking(move || {
                std::net::ToSocketAddrs::to_socket_addrs(&s)
            }))
        }
    }

    // ===== impl (&str, u16) =====

    impl ToSocketAddrs for (&str, u16) {}

    impl sealed::ToSocketAddrsPriv for (&str, u16) {
        type Iter = sealed::OneOrMore;
        type Future = sealed::MaybeReady;

        fn to_socket_addrs(&self) -> Self::Future {
            use crate::runtime::spawn_blocking;
            use sealed::MaybeReady;

            let (host, port) = *self;

            // try to parse the host as a regular IP address first
            if let Ok(addr) = host.parse::<Ipv4Addr>() {
                let addr = SocketAddrV4::new(addr, port);
                let addr = SocketAddr::V4(addr);

                return MaybeReady::Ready(Some(addr));
            }

            if let Ok(addr) = host.parse::<Ipv6Addr>() {
                let addr = SocketAddrV6::new(addr, port, 0, 0);
                let addr = SocketAddr::V6(addr);

                return MaybeReady::Ready(Some(addr));
            }

            let host = host.to_owned();

            MaybeReady::Blocking(spawn_blocking(move || {
                std::net::ToSocketAddrs::to_socket_addrs(&(&host[..], port))
            }))
        }
    }

    // ===== impl String =====

    impl ToSocketAddrs for String {}

    impl sealed::ToSocketAddrsPriv for String {
        type Iter = <str as sealed::ToSocketAddrsPriv>::Iter;
        type Future = <str as sealed::ToSocketAddrsPriv>::Future;

        fn to_socket_addrs(&self) -> Self::Future {
            (&self[..]).to_socket_addrs()
        }
    }
}

pub(crate) mod sealed {
    //! The contents of this trait are intended to remain private and __not__
    //! part of the `ToSocketAddrs` public API. The details will change over
    //! time.

    use std::future::Future;
    use std::io;
    use std::net::SocketAddr;

    cfg_dns! {
        use crate::task::JoinHandle;

        use std::option;
        use std::pin::Pin;
        use std::task::{Context, Poll};
        use std::vec;
    }

    #[doc(hidden)]
    pub trait ToSocketAddrsPriv {
        type Iter: Iterator<Item = SocketAddr> + Send + 'static;
        type Future: Future<Output = io::Result<Self::Iter>> + Send + 'static;

        fn to_socket_addrs(&self) -> Self::Future;
    }

    cfg_dns! {
        #[doc(hidden)]
        #[derive(Debug)]
        pub enum MaybeReady {
            Ready(Option<SocketAddr>),
            Blocking(JoinHandle<io::Result<vec::IntoIter<SocketAddr>>>),
        }

        #[doc(hidden)]
        #[derive(Debug)]
        pub enum OneOrMore {
            One(option::IntoIter<SocketAddr>),
            More(vec::IntoIter<SocketAddr>),
        }

        impl Future for MaybeReady {
            type Output = io::Result<OneOrMore>;

            fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                match *self {
                    MaybeReady::Ready(ref mut i) => {
                        let iter = OneOrMore::One(i.take().into_iter());
                        Poll::Ready(Ok(iter))
                    }
                    MaybeReady::Blocking(ref mut rx) => {
                        let res = ready!(Pin::new(rx).poll(cx))?.map(OneOrMore::More);

                        Poll::Ready(res)
                    }
                }
            }
        }

        impl Iterator for OneOrMore {
            type Item = SocketAddr;

            fn next(&mut self) -> Option<Self::Item> {
                match self {
                    OneOrMore::One(i) => i.next(),
                    OneOrMore::More(i) => i.next(),
                }
            }

            fn size_hint(&self) -> (usize, Option<usize>) {
                match self {
                    OneOrMore::One(i) => i.size_hint(),
                    OneOrMore::More(i) => i.size_hint(),
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use request::Request;
use response::Response;
use nickel_error::NickelError;
use hyper::net;

pub use self::Action::{Continue, Halt};

pub type MiddlewareResult<'mw, D> = Result<Action<Response<'mw, D, net::Fresh>,
                                              Response<'mw, D, net::Streaming>>,
                                        NickelError<'mw, D>>;

pub enum Action<T=(), U=()> {
    Continue(T),
    Halt(U)
}

// the usage of + Send is weird here because what we really want is + Static
// but that's not possible as of today. We have to use + Send for now.
pub trait Middleware<D>: Send + 'static + Sync {
    fn invoke<'mw, 'conn>(&'mw self, _req: &mut Request<'mw, 'conn, D>, res: Response<'mw, D, net::Fresh>) -> MiddlewareResult<'mw, D> {
        Ok(Continue(res))
    }
}

impl<T, D> Middleware<D> for T where T: for<'r, 'mw, 'conn> Fn(&'r mut Request<'mw, 'conn, D>, Response<'mw, D>) -> MiddlewareResult<'mw, D> + Send + Sync + 'static {
    fn invoke<'mw, 'conn>(&'mw self, req: &mut Request<'mw, 'conn, D>, res: Response<'mw, D>) -> MiddlewareResult<'mw, D> {
        (*self)(req, res)
    }
}

pub trait ErrorHandler<D>: Send + 'static + Sync {
    fn handle_error(&self, &mut NickelError<D>, &mut Request<D>) -> Action;
}

impl<D: 'static> ErrorHandler<D> for fn(&mut NickelError<D>, &mut Request<D>) -> Action {
    fn handle_error(&self, err: &mut NickelError<D>, req: &mut Request<D>) -> Action {
        (*self)(err, req)
    }
}

pub struct MiddlewareStack<D=()> {
    handlers: Vec<Box<Middleware<D> + Send + Sync>>,
    error_handlers: Vec<Box<ErrorHandler<D> + Send + Sync>>
}

impl<D: 'static> MiddlewareStack<D> {
    pub fn add_middleware<T: Middleware<D>> (&mut self, handler: T) {
        self.handlers.push(Box::new(handler));
    }

    pub fn add_error_handler<T: ErrorHandler<D>> (&mut self, handler: T) {
        self.error_handlers.push(Box::new(handler));
    }

    pub fn invoke<'mw, 'conn>(&'mw self, mut req: Request<'mw, 'conn, D>, mut res: Response<'mw, D>) {
        for handler in self.handlers.iter() {
            match handler.invoke(&mut req, res) {
                Ok(Halt(res)) => {
                    debug!("Halted {:?} {:?} {:?} {:?}",
                           req.origin.method,
                           req.origin.remote_addr,
                           req.origin.uri,
                           res.status());
                    let _ = res.end();
                    return
                }
                Ok(Continue(fresh)) => res = fresh,
                Err(mut err) => {
                    warn!("{:?} {:?} {:?} {:?} {:?}",
                          req.origin.method,
                          req.origin.remote_addr,
                          req.origin.uri,
                          err.message,
                          err.stream.as_ref().map(|s| s.status()));

                    for error_handler in self.error_handlers.iter().rev() {
                        if let Halt(()) = error_handler.handle_error(&mut err, &mut req) {
                            err.end();
                            return
                        }
                    }

                    warn!("Unhandled error: {:?} {:?} {:?} {:?} {:?}",
                          req.origin.method,
                          req.origin.remote_addr,
                          req.origin.uri,
                          err.message,
                          err.stream.map(|s| s.status()));
                    return
                }
            }
        }
    }

    pub fn new () -> MiddlewareStack<D> {
        MiddlewareStack{
            handlers: Vec::new(),
            error_handlers: Vec::new()
        }
    }
}
<commit_msg>fix(*): add default typeparam for MiddlewareResult<commit_after>use request::Request;
use response::Response;
use nickel_error::NickelError;
use hyper::net;

pub use self::Action::{Continue, Halt};

pub type MiddlewareResult<'mw, D= ()> = Result<Action<Response<'mw, D, net::Fresh>,
                                                      Response<'mw, D, net::Streaming>>,
                                               NickelError<'mw, D>>;

pub enum Action<T=(), U=()> {
    Continue(T),
    Halt(U)
}

// the usage of + Send is weird here because what we really want is + Static
// but that's not possible as of today. We have to use + Send for now.
pub trait Middleware<D>: Send + 'static + Sync {
    fn invoke<'mw, 'conn>(&'mw self, _req: &mut Request<'mw, 'conn, D>, res: Response<'mw, D, net::Fresh>) -> MiddlewareResult<'mw, D> {
        Ok(Continue(res))
    }
}

impl<T, D> Middleware<D> for T where T: for<'r, 'mw, 'conn> Fn(&'r mut Request<'mw, 'conn, D>, Response<'mw, D>) -> MiddlewareResult<'mw, D> + Send + Sync + 'static {
    fn invoke<'mw, 'conn>(&'mw self, req: &mut Request<'mw, 'conn, D>, res: Response<'mw, D>) -> MiddlewareResult<'mw, D> {
        (*self)(req, res)
    }
}

pub trait ErrorHandler<D>: Send + 'static + Sync {
    fn handle_error(&self, &mut NickelError<D>, &mut Request<D>) -> Action;
}

impl<D: 'static> ErrorHandler<D> for fn(&mut NickelError<D>, &mut Request<D>) -> Action {
    fn handle_error(&self, err: &mut NickelError<D>, req: &mut Request<D>) -> Action {
        (*self)(err, req)
    }
}

pub struct MiddlewareStack<D=()> {
    handlers: Vec<Box<Middleware<D> + Send + Sync>>,
    error_handlers: Vec<Box<ErrorHandler<D> + Send + Sync>>
}

impl<D: 'static> MiddlewareStack<D> {
    pub fn add_middleware<T: Middleware<D>> (&mut self, handler: T) {
        self.handlers.push(Box::new(handler));
    }

    pub fn add_error_handler<T: ErrorHandler<D>> (&mut self, handler: T) {
        self.error_handlers.push(Box::new(handler));
    }

    pub fn invoke<'mw, 'conn>(&'mw self, mut req: Request<'mw, 'conn, D>, mut res: Response<'mw, D>) {
        for handler in self.handlers.iter() {
            match handler.invoke(&mut req, res) {
                Ok(Halt(res)) => {
                    debug!("Halted {:?} {:?} {:?} {:?}",
                           req.origin.method,
                           req.origin.remote_addr,
                           req.origin.uri,
                           res.status());
                    let _ = res.end();
                    return
                }
                Ok(Continue(fresh)) => res = fresh,
                Err(mut err) => {
                    warn!("{:?} {:?} {:?} {:?} {:?}",
                          req.origin.method,
                          req.origin.remote_addr,
                          req.origin.uri,
                          err.message,
                          err.stream.as_ref().map(|s| s.status()));

                    for error_handler in self.error_handlers.iter().rev() {
                        if let Halt(()) = error_handler.handle_error(&mut err, &mut req) {
                            err.end();
                            return
                        }
                    }

                    warn!("Unhandled error: {:?} {:?} {:?} {:?} {:?}",
                          req.origin.method,
                          req.origin.remote_addr,
                          req.origin.uri,
                          err.message,
                          err.stream.map(|s| s.status()));
                    return
                }
            }
        }
    }

    pub fn new () -> MiddlewareStack<D> {
        MiddlewareStack{
            handlers: Vec::new(),
            error_handlers: Vec::new()
        }
    }
}
<|endoftext|>