<commit_before>use crate::{cluster::Cluster, instance::Instance};
use admission_control_proto::proto::{
    admission_control::SubmitTransactionRequest, admission_control_grpc::AdmissionControlClient,
};
use grpcio::{ChannelBuilder, EnvBuilder};
use proto_conv::IntoProto;
use std::{
    env, slice,
    str::FromStr,
    sync::Arc,
    thread,
    time::{Duration, Instant},
};

use crate::{prometheus::Prometheus, util::unix_timestamp_now};
use admission_control_proto::{AdmissionControlStatus, SubmitTransactionResponse};
use crypto::{
    ed25519::{Ed25519PrivateKey, Ed25519PublicKey},
    test_utils::KeyPair,
    traits::Uniform,
};
use failure::{
    self,
    prelude::{bail, format_err},
};
use generate_keypair::load_key_from_file;
use itertools::zip;
use proto_conv::FromProto;
use rand::{
    rngs::{EntropyRng, StdRng},
    Rng, SeedableRng,
};
use types::{
    account_address::AccountAddress,
    account_config::{association_address, get_account_resource_or_default},
    get_with_proof::ResponseItem,
    proto::get_with_proof::{
        GetAccountStateRequest, RequestItem, RequestItem_oneof_requested_items,
        UpdateToLatestLedgerRequest,
    },
    transaction::{Script, TransactionPayload},
    transaction_helpers::create_signed_txn,
};

const ACCOUNT_PER_CLIENT_DEFAULT: usize = 10;
const THREADS_PER_CLIENT_DEFAULT: usize = 1;
const MAX_TXN_BATCH_SIZE: usize = 100; // Max transactions per account in mempool

pub struct TxEmitter {
    clients: Vec<(Instance, AdmissionControlClient)>,
}

impl TxEmitter {
    pub fn new(cluster: &Cluster) -> Self {
        let clients = Self::create_ac_clients(cluster);

        Self { clients }
    }

    fn create_ac_clients(cluster: &Cluster) -> Vec<(Instance, AdmissionControlClient)> {
        let mut clients = vec![];
        let threads_per_client = get_env("THREADS_PER_CLIENT", THREADS_PER_CLIENT_DEFAULT);
        for instance in cluster.instances() {
            let address = format!("{}:8000", instance.ip());
            for _ in 0..threads_per_client {
                let env_builder = Arc::new(EnvBuilder::new().name_prefix("ac-grpc-").build());
                let ch = ChannelBuilder::new(env_builder).connect(&address);
                clients.push((instance.clone(), AdmissionControlClient::new(ch)));
            }
        }
        clients
    }

    pub fn run(self) {
        let mint_client = &self.clients[0].1;
        let mut faucet_account = load_faucet_account(mint_client, "mint.key");
        let account_per_client = get_env("ACCOUNT_PER_CLIENT", ACCOUNT_PER_CLIENT_DEFAULT);
        let num_accounts = account_per_client * self.clients.len();
        println!("Minting accounts");
        let mut all_accounts: Vec<AccountData> = vec![];
        for _ in 0..(num_accounts + MAX_TXN_BATCH_SIZE - 1) / MAX_TXN_BATCH_SIZE {
            let mut accounts = gen_random_accounts(MAX_TXN_BATCH_SIZE);
            let mint_requests = gen_mint_txn_requests(&mut faucet_account, &accounts);
            execute_and_wait_transactions(&mint_client, &faucet_account, mint_requests);
            all_accounts.append(&mut accounts);
        }
        println!("Mint is done");
        let mut join_handles = vec![];
        let mut all_accounts = all_accounts.into_iter();
        for (index, (instance, client)) in self.clients.into_iter().enumerate() {
            let accounts = (&mut all_accounts).take(account_per_client).collect();
            let thread = SubmissionThread {
                accounts,
                instance,
                client,
            };
            let join_handle = thread::Builder::new()
                .name(format!("thread-{}", index))
                .spawn(move || thread.run())
                .unwrap();
            join_handles.push(join_handle);
            thread::sleep(Duration::from_millis(10)); // Small stagger between starting threads
        }
        println!("Threads started");
        let prometheus = Prometheus::try_new_from_environment();
        if let Some(prometheus) = prometheus {
            run_stat_loop(prometheus);
        }
        for join_handle in join_handles {
            join_handle.join().unwrap();
        }
    }
}

fn run_stat_loop(prometheus: Prometheus) {
    thread::sleep(Duration::from_secs(30)); // warm up
    loop {
        thread::sleep(Duration::from_secs(10));
        if let Err(err) = print_stat(&prometheus) {
            println!("Stat error: {:?}", err);
        }
    }
}

fn print_stat(prometheus: &Prometheus) -> failure::Result<()> {
    let step = 10;
    let end = unix_timestamp_now();
    let start = end - Duration::from_secs(30); // avg over last 30 sec
    let tps = prometheus.query_range(
        "irate(consensus_gauge{op='last_committed_version'}[1m])".to_string(),
        &start,
        &end,
        step,
    )?;
    let avg_tps = tps.avg().ok_or_else(|| format_err!("No tps data"))?;
    let latency = prometheus.query_range(
        "irate(mempool_duration_sum{op='e2e.latency'}[1m])/irate(mempool_duration_count{op='e2e.latency'}[1m])"
            .to_string(),
        &start,
        &end,
        step,
    )?;
    let avg_latency = latency
        .avg()
        .ok_or_else(|| format_err!("No latency data"))?;
    println!(
        "Tps: {:.0}, latency: {:.0} ms",
        avg_tps,
        avg_latency * 1000.
    );
    Ok(())
}

fn get_env<F: FromStr>(name: &str, default: F) -> F {
    match env::var(name) {
        Ok(v) => match v.parse() {
            Ok(v) => v,
            _ => panic!("Failed to parse env {}", name),
        },
        _ => default,
    }
}

struct SubmissionThread {
    accounts: Vec<AccountData>,
    instance: Instance,
    client: AdmissionControlClient,
}

impl SubmissionThread {
    fn run(mut self) {
        let wait_millis = get_env("WAIT_MILLIS", 50);
        let wait = Duration::from_millis(wait_millis);
        let wait_committed = get_env("WAIT_COMMITTED", true);
        let verbose = get_env("VERBOSE", false);
        loop {
            let gen_requests = gen_ring_requests(&mut self.accounts);
            for request in gen_requests {
                let wait_util = Instant::now() + wait;
                let resp = self.client.submit_transaction(&request);
                match resp {
                    Err(e) => {
                        if verbose {
                            println!("Failed to submit request to {}: {:?}", self.instance, e);
                        }
                    }
                    Ok(r) => {
                        let r = SubmitTransactionResponse::from_proto(r)
                            .expect("Failed to parse SubmitTransactionResponse");
                        if verbose && !is_accepted(&r) {
                            println!("Request declined: {:?}", r);
                        }
                    }
                }
                let now = Instant::now();
                if wait_util > now {
                    thread::sleep(wait_util - now);
                } else if verbose {
                    println!("Thread for {} won't sleep", self.instance);
                }
            }
            if wait_committed {
                wait_for_accounts_sequence(&self.client, &self.accounts);
            }
        }
    }
}

fn wait_for_accounts_sequence(client: &AdmissionControlClient, accounts: &[AccountData]) {
    let addresses: Vec<_> = accounts.iter().map(|d| d.address).collect();
    loop {
        match query_sequence_numbers(client, &addresses) {
            Err(e) => println!("Failed to query ledger info: {:?}", e),
            Ok(sequence_numbers) => {
                if is_sequence_equal(accounts, &sequence_numbers) {
                    break;
                }
            }
        }
        thread::sleep(Duration::from_millis(100));
    }
}

fn is_sequence_equal(accounts: &[AccountData], sequence_numbers: &[u64]) -> bool {
    for (account, sequence_number) in zip(accounts, sequence_numbers) {
        if *sequence_number != account.sequence_number {
            return false;
        }
    }
    true
}

fn query_sequence_numbers(
    client: &AdmissionControlClient,
    addresses: &[AccountAddress],
) -> failure::Result<Vec<u64>> {
    let mut update_request = UpdateToLatestLedgerRequest::new();
    for address in addresses {
        let mut request_item = RequestItem::new();
        let mut account_state_request = GetAccountStateRequest::new();
        account_state_request.address = address.to_vec();
        request_item.requested_items = Some(
            RequestItem_oneof_requested_items::get_account_state_request(account_state_request),
        );
        update_request.requested_items.push(request_item);
    }
    let resp = client
        .update_to_latest_ledger(&update_request)
        .map_err(|e| format_err!("update_to_latest_ledger failed: {:?} ", e))?;
    let mut result = Vec::with_capacity(resp.response_items.len());
    for item in resp.response_items.into_iter() {
        let item = ResponseItem::from_proto(item)
            .map_err(|e| format_err!("ResponseItem::from_proto failed: {:?} ", e))?;
        if let ResponseItem::GetAccountState {
            account_state_with_proof,
        } = item
        {
            let account_resource = get_account_resource_or_default(&account_state_with_proof.blob)
                .map_err(|e| format_err!("get_account_resource_or_default failed: {:?} ", e))?;
            result.push(account_resource.sequence_number());
        } else {
            bail!(
                "Unexpected item in UpdateToLatestLedgerResponse: {:?}",
                item
            );
        }
    }
    Ok(result)
}

const MAX_GAS_AMOUNT: u64 = 1_000_000;
const GAS_UNIT_PRICE: u64 = 0;
const TXN_EXPIRATION: i64 = 100;

fn gen_submit_transaction_request(
    script: Script,
    sender_account: &mut AccountData,
) -> SubmitTransactionRequest {
    let signed_txn = create_signed_txn(
        &sender_account.key_pair,
        TransactionPayload::Script(script),
        sender_account.address,
        sender_account.sequence_number,
        MAX_GAS_AMOUNT,
        GAS_UNIT_PRICE,
        TXN_EXPIRATION,
    )
    .expect("Failed to create signed transaction");
    let mut req = SubmitTransactionRequest::new();
    req.set_signed_txn(signed_txn.into_proto());
    sender_account.sequence_number += 1;
    req
}

fn gen_transfer_txn_request(
    sender: &mut AccountData,
    receiver: &AccountAddress,
    num_coins: u64,
) -> SubmitTransactionRequest {
    let script = transaction_builder::encode_transfer_script(&receiver, num_coins);
    gen_submit_transaction_request(script, sender)
}

fn gen_ring_requests(accounts: &mut [AccountData]) -> Vec<SubmitTransactionRequest> {
    let mut receiver_addrs: Vec<AccountAddress> =
        accounts.iter().map(|account| account.address).collect();
    receiver_addrs.rotate_left(1);
    accounts
        .iter_mut()
        .zip(receiver_addrs.iter())
        .map(|(sender, receiver_addr)| gen_transfer_txn_request(sender, receiver_addr, 1))
        .collect()
}

fn gen_random_account(rng: &mut StdRng) -> AccountData {
    let key_pair = KeyPair::generate_for_testing(rng);
    AccountData {
        address: AccountAddress::from_public_key(&key_pair.public_key),
        key_pair,
        sequence_number: 0,
    }
}

fn gen_random_accounts(num_accounts: usize) -> Vec<AccountData> {
    let seed: [u8; 32] = EntropyRng::new().gen();
    let mut rng = StdRng::from_seed(seed);
    (0..num_accounts)
        .map(|_| gen_random_account(&mut rng))
        .collect()
}

fn gen_mint_txn_request(
    faucet_account: &mut AccountData,
    receiver: &AccountAddress,
) -> SubmitTransactionRequest {
    let program = transaction_builder::encode_mint_script(receiver, 1_000_000);
    gen_submit_transaction_request(program, faucet_account)
}

fn gen_mint_txn_requests(
    faucet_account: &mut AccountData,
    accounts: &[AccountData],
) -> Vec<SubmitTransactionRequest> {
    accounts
        .iter()
        .map(|account| gen_mint_txn_request(faucet_account, &account.address))
        .collect()
}

fn execute_and_wait_transactions(
    client: &AdmissionControlClient,
    account: &AccountData,
    txn: Vec<SubmitTransactionRequest>,
) {
    for request in txn {
        let resp = client.submit_transaction(&request);
        match resp {
            Err(e) => println!("Failed to submit request: {:?}", e),
            Ok(r) => {
                let r = SubmitTransactionResponse::from_proto(r)
                    .expect("Failed to parse SubmitTransactionResponse");
                if !is_accepted(&r) {
                    println!("Request declined: {:?}", r);
                }
            }
        }
    }
    wait_for_accounts_sequence(client, slice::from_ref(account));
}

fn load_faucet_account(client: &AdmissionControlClient, faucet_account_path: &str) -> AccountData {
    let key_pair: KeyPair<Ed25519PrivateKey, Ed25519PublicKey> =
        load_key_from_file(faucet_account_path).expect("invalid faucet keypair file");
    let address = association_address();
    let sequence_number = query_sequence_numbers(client, &[address])
        .expect("query_sequence_numbers for faucet account failed")[0];
    AccountData {
        address,
        key_pair,
        sequence_number,
    }
}

struct AccountData {
    pub address: AccountAddress,
    pub key_pair: KeyPair<Ed25519PrivateKey, Ed25519PublicKey>,
    pub sequence_number: u64,
}

fn is_accepted(resp: &SubmitTransactionResponse) -> bool {
    if let Some(ref status) = resp.ac_status {
        return *status == AdmissionControlStatus::Accepted;
    }
    false
}
<commit_msg>[cluster_test] Send request to random account during transaction generation<commit_after>use crate::{cluster::Cluster, instance::Instance};
use admission_control_proto::proto::{
    admission_control::SubmitTransactionRequest, admission_control_grpc::AdmissionControlClient,
};
use grpcio::{ChannelBuilder, EnvBuilder};
use proto_conv::IntoProto;
use std::{
    env, slice,
    str::FromStr,
    sync::Arc,
    thread,
    time::{Duration, Instant},
};

use crate::{prometheus::Prometheus, util::unix_timestamp_now};
use admission_control_proto::{AdmissionControlStatus, SubmitTransactionResponse};
use crypto::{
    ed25519::{Ed25519PrivateKey, Ed25519PublicKey},
    test_utils::KeyPair,
    traits::Uniform,
};
use failure::{
    self,
    prelude::{bail, format_err},
};
use generate_keypair::load_key_from_file;
use itertools::zip;
use proto_conv::FromProto;
use rand::{
    prelude::ThreadRng,
    rngs::{EntropyRng, StdRng},
    seq::SliceRandom,
    Rng, SeedableRng,
};
use types::{
    account_address::AccountAddress,
    account_config::{association_address, get_account_resource_or_default},
    get_with_proof::ResponseItem,
    proto::get_with_proof::{
        GetAccountStateRequest, RequestItem, RequestItem_oneof_requested_items,
        UpdateToLatestLedgerRequest,
    },
    transaction::{Script, TransactionPayload},
    transaction_helpers::create_signed_txn,
};

const ACCOUNT_PER_CLIENT_DEFAULT: usize = 10;
const THREADS_PER_CLIENT_DEFAULT: usize = 1;
const MAX_TXN_BATCH_SIZE: usize = 100; // Max transactions per account in mempool

pub struct TxEmitter {
    clients: Vec<(Instance, AdmissionControlClient)>,
}

impl TxEmitter {
    pub fn new(cluster: &Cluster) -> Self {
        let clients = Self::create_ac_clients(cluster);

        Self { clients }
    }

    fn create_ac_clients(cluster: &Cluster) -> Vec<(Instance, AdmissionControlClient)> {
        let mut clients = vec![];
        let threads_per_client = get_env("THREADS_PER_CLIENT", THREADS_PER_CLIENT_DEFAULT);
        for instance in cluster.instances() {
            let address = format!("{}:8000", instance.ip());
            for _ in 0..threads_per_client {
                let env_builder = Arc::new(EnvBuilder::new().name_prefix("ac-grpc-").build());
                let ch = ChannelBuilder::new(env_builder).connect(&address);
                clients.push((instance.clone(), AdmissionControlClient::new(ch)));
            }
        }
        clients
    }

    pub fn run(self) {
        let mint_client = &self.clients[0].1;
        let mut faucet_account = load_faucet_account(mint_client, "mint.key");
        let account_per_client = get_env("ACCOUNT_PER_CLIENT", ACCOUNT_PER_CLIENT_DEFAULT);
        let num_accounts = account_per_client * self.clients.len();
        println!("Minting accounts");
        let mut all_accounts: Vec<AccountData> = vec![];
        for _ in 0..(num_accounts + MAX_TXN_BATCH_SIZE - 1) / MAX_TXN_BATCH_SIZE {
            let mut accounts = gen_random_accounts(MAX_TXN_BATCH_SIZE);
            let mint_requests = gen_mint_txn_requests(&mut faucet_account, &accounts);
            execute_and_wait_transactions(&mint_client, &faucet_account, mint_requests);
            all_accounts.append(&mut accounts);
        }
        println!("Mint is done");
        let mut join_handles = vec![];
        let all_addresses: Vec<_> = all_accounts.iter().map(|d| d.address).collect();
        let all_addresses = Arc::new(all_addresses);
        let mut all_accounts = all_accounts.into_iter();
        for (index, (instance, client)) in self.clients.into_iter().enumerate() {
            let accounts = (&mut all_accounts).take(account_per_client).collect();
            let all_addresses = all_addresses.clone();
            let thread = SubmissionThread {
                accounts,
                instance,
                client,
                all_addresses,
            };
            let join_handle = thread::Builder::new()
                .name(format!("thread-{}", index))
                .spawn(move || thread.run())
                .unwrap();
            join_handles.push(join_handle);
            thread::sleep(Duration::from_millis(10)); // Small stagger between starting threads
        }
        println!("Threads started");
        let prometheus = Prometheus::try_new_from_environment();
        if let Some(prometheus) = prometheus {
            run_stat_loop(prometheus);
        }
        for join_handle in join_handles {
            join_handle.join().unwrap();
        }
    }
}

fn run_stat_loop(prometheus: Prometheus) {
    thread::sleep(Duration::from_secs(30)); // warm up
    loop {
        thread::sleep(Duration::from_secs(10));
        if let Err(err) = print_stat(&prometheus) {
            println!("Stat error: {:?}", err);
        }
    }
}

fn print_stat(prometheus: &Prometheus) -> failure::Result<()> {
    let step = 10;
    let end = unix_timestamp_now();
    let start = end - Duration::from_secs(30); // avg over last 30 sec
    let tps = prometheus.query_range(
        "irate(consensus_gauge{op='last_committed_version'}[1m])".to_string(),
        &start,
        &end,
        step,
    )?;
    let avg_tps = tps.avg().ok_or_else(|| format_err!("No tps data"))?;
    let latency = prometheus.query_range(
        "irate(mempool_duration_sum{op='e2e.latency'}[1m])/irate(mempool_duration_count{op='e2e.latency'}[1m])"
            .to_string(),
        &start,
        &end,
        step,
    )?;
    let avg_latency = latency
        .avg()
        .ok_or_else(|| format_err!("No latency data"))?;
    println!(
        "Tps: {:.0}, latency: {:.0} ms",
        avg_tps,
        avg_latency * 1000.
    );
    Ok(())
}

fn get_env<F: FromStr>(name: &str, default: F) -> F {
    match env::var(name) {
        Ok(v) => match v.parse() {
            Ok(v) => v,
            _ => panic!("Failed to parse env {}", name),
        },
        _ => default,
    }
}

struct SubmissionThread {
    accounts: Vec<AccountData>,
    instance: Instance,
    client: AdmissionControlClient,
    all_addresses: Arc<Vec<AccountAddress>>,
}

impl SubmissionThread {
    fn run(mut self) {
        let wait_millis = get_env("WAIT_MILLIS", 50);
        let wait = Duration::from_millis(wait_millis);
        let wait_committed = get_env("WAIT_COMMITTED", true);
        let verbose = get_env("VERBOSE", false);
        let mut rng = ThreadRng::default();
        loop {
            let requests = self.gen_requests(&mut rng);
            for request in requests {
                let wait_util = Instant::now() + wait;
                let resp = self.client.submit_transaction(&request);
                match resp {
                    Err(e) => {
                        if verbose {
                            println!("Failed to submit request to {}: {:?}", self.instance, e);
                        }
                    }
                    Ok(r) => {
                        let r = SubmitTransactionResponse::from_proto(r)
                            .expect("Failed to parse SubmitTransactionResponse");
                        if verbose && !is_accepted(&r) {
                            println!("Request declined: {:?}", r);
                        }
                    }
                }
                let now = Instant::now();
                if wait_util > now {
                    thread::sleep(wait_util - now);
                } else if verbose {
                    println!("Thread for {} won't sleep", self.instance);
                }
            }
            if wait_committed {
                wait_for_accounts_sequence(&self.client, &self.accounts);
            }
        }
    }

    fn gen_requests(&mut self, rng: &mut ThreadRng) -> Vec<SubmitTransactionRequest> {
        let mut requests = Vec::with_capacity(self.accounts.len());
        let all_addresses = &self.all_addresses;
        for sender in &mut self.accounts {
            let receiver = all_addresses
                .choose(rng)
                .expect("all_addresses can't be empty");
            let request = gen_transfer_txn_request(sender, receiver, 1);
            requests.push(request);
        }
        requests
    }
}

fn wait_for_accounts_sequence(client: &AdmissionControlClient, accounts: &[AccountData]) {
    let addresses: Vec<_> = accounts.iter().map(|d| d.address).collect();
    loop {
        match query_sequence_numbers(client, &addresses) {
            Err(e) => println!("Failed to query ledger info: {:?}", e),
            Ok(sequence_numbers) => {
                if is_sequence_equal(accounts, &sequence_numbers) {
                    break;
                }
            }
        }
        thread::sleep(Duration::from_millis(100));
    }
}

fn is_sequence_equal(accounts: &[AccountData], sequence_numbers: &[u64]) -> bool {
    for (account, sequence_number) in zip(accounts, sequence_numbers) {
        if *sequence_number != account.sequence_number {
            return false;
        }
    }
    true
}

fn query_sequence_numbers(
    client: &AdmissionControlClient,
    addresses: &[AccountAddress],
) -> failure::Result<Vec<u64>> {
    let mut update_request = UpdateToLatestLedgerRequest::new();
    for address in addresses {
        let mut request_item = RequestItem::new();
        let mut account_state_request = GetAccountStateRequest::new();
        account_state_request.address = address.to_vec();
        request_item.requested_items = Some(
            RequestItem_oneof_requested_items::get_account_state_request(account_state_request),
        );
        update_request.requested_items.push(request_item);
    }
    let resp = client
        .update_to_latest_ledger(&update_request)
        .map_err(|e| format_err!("update_to_latest_ledger failed: {:?} ", e))?;
    let mut result = Vec::with_capacity(resp.response_items.len());
    for item in resp.response_items.into_iter() {
        let item = ResponseItem::from_proto(item)
            .map_err(|e| format_err!("ResponseItem::from_proto failed: {:?} ", e))?;
        if let ResponseItem::GetAccountState {
            account_state_with_proof,
        } = item
        {
            let account_resource = get_account_resource_or_default(&account_state_with_proof.blob)
                .map_err(|e| format_err!("get_account_resource_or_default failed: {:?} ", e))?;
            result.push(account_resource.sequence_number());
        } else {
            bail!(
                "Unexpected item in UpdateToLatestLedgerResponse: {:?}",
                item
            );
        }
    }
    Ok(result)
}

const MAX_GAS_AMOUNT: u64 = 1_000_000;
const GAS_UNIT_PRICE: u64 = 0;
const TXN_EXPIRATION: i64 = 100;

fn gen_submit_transaction_request(
    script: Script,
    sender_account: &mut AccountData,
) -> SubmitTransactionRequest {
    let signed_txn = create_signed_txn(
        &sender_account.key_pair,
        TransactionPayload::Script(script),
        sender_account.address,
        sender_account.sequence_number,
        MAX_GAS_AMOUNT,
        GAS_UNIT_PRICE,
        TXN_EXPIRATION,
    )
    .expect("Failed to create signed transaction");
    let mut req = SubmitTransactionRequest::new();
    req.set_signed_txn(signed_txn.into_proto());
    sender_account.sequence_number += 1;
    req
}

fn gen_transfer_txn_request(
    sender: &mut AccountData,
    receiver: &AccountAddress,
    num_coins: u64,
) -> SubmitTransactionRequest {
    let script = transaction_builder::encode_transfer_script(&receiver, num_coins);
    gen_submit_transaction_request(script, sender)
}

fn gen_random_account(rng: &mut StdRng) -> AccountData {
    let key_pair = KeyPair::generate_for_testing(rng);
    AccountData {
        address: AccountAddress::from_public_key(&key_pair.public_key),
        key_pair,
        sequence_number: 0,
    }
}

fn gen_random_accounts(num_accounts: usize) -> Vec<AccountData> {
    let seed: [u8; 32] = EntropyRng::new().gen();
    let mut rng = StdRng::from_seed(seed);
    (0..num_accounts)
        .map(|_| gen_random_account(&mut rng))
        .collect()
}

fn gen_mint_txn_request(
    faucet_account: &mut AccountData,
    receiver: &AccountAddress,
) -> SubmitTransactionRequest {
    let program = transaction_builder::encode_mint_script(receiver, 1_000_000);
    gen_submit_transaction_request(program, faucet_account)
}

fn gen_mint_txn_requests(
    faucet_account: &mut AccountData,
    accounts: &[AccountData],
) -> Vec<SubmitTransactionRequest> {
    accounts
        .iter()
        .map(|account| gen_mint_txn_request(faucet_account, &account.address))
        .collect()
}

fn execute_and_wait_transactions(
    client: &AdmissionControlClient,
    account: &AccountData,
    txn: Vec<SubmitTransactionRequest>,
) {
    for request in txn {
        let resp = client.submit_transaction(&request);
        match resp {
            Err(e) => println!("Failed to submit request: {:?}", e),
            Ok(r) => {
                let r = SubmitTransactionResponse::from_proto(r)
                    .expect("Failed to parse SubmitTransactionResponse");
                if !is_accepted(&r) {
                    println!("Request declined: {:?}", r);
                }
            }
        }
    }
    wait_for_accounts_sequence(client, slice::from_ref(account));
}

fn load_faucet_account(client: &AdmissionControlClient, faucet_account_path: &str) -> AccountData {
    let key_pair: KeyPair<Ed25519PrivateKey, Ed25519PublicKey> =
        load_key_from_file(faucet_account_path).expect("invalid faucet keypair file");
    let address = association_address();
    let sequence_number = query_sequence_numbers(client, &[address])
        .expect("query_sequence_numbers for faucet account failed")[0];
    AccountData {
        address,
        key_pair,
        sequence_number,
    }
}

struct AccountData {
    pub address: AccountAddress,
    pub key_pair: KeyPair<Ed25519PrivateKey, Ed25519PublicKey>,
    pub sequence_number: u64,
}

fn is_accepted(resp: &SubmitTransactionResponse) -> bool {
    if let Some(ref status) = resp.ac_status {
        return *status == AdmissionControlStatus::Accepted;
    }
    false
}
<|endoftext|>


<commit_before>use json_flex;
use json_flex::JFObject;

#[derive(Clone)]
pub struct Rows {
    data: Option<Vec<JFObject>>
}

/// A result rows representing type.
/// This type is usually generated by `ResultParser#into_raw()`.
impl Rows {
    pub fn new(data: Option<Vec<JFObject>>) -> Rows {
        Rows { data: data }
    }

    /// Get columns in `Rows`.
    pub fn columns(&mut self) -> Option<Vec<JFObject>> {
        let popable = match self.data.clone() {
            Some(v) => v,
            None => return None
        };
        let pop = match popable.clone().pop() {
            Some(v) => v,
            None => return None
        };
        match pop.into_vec().clone() {
            Some(v) => return Some(v.clone()),
            None => return None
        }
    }
}

#[derive(Clone, Debug)]
pub struct ResultParser {
    result: Box<JFObject>
}

/// ResultParser
///
/// Groonga returns following array json:
///
/// success response:
/// `[[status,start_time, elapsed_time],
///   [[[matched_columns],[Array([column1,type1], ...)],
///     [Array([result1, result2, ...])]]]]`
///
/// error response:
/// `[[status,start_time, elapsed_time, error_information, ...]]`
impl ResultParser {
    pub fn new(json: String) -> ResultParser {
        ResultParser{result: json_flex::decode(json)}
    }

    /// Get raw response result.
    pub fn get_raw_object(&mut self) -> Box<JFObject> {
        self.result.clone()
    }

    /// Return header elements in response.
    ///
    /// # Panics
    ///
    /// Panics if response json is corrupted.
    pub fn get_header(&mut self) -> JFObject {
        self.result[0].clone()
    }

    /// Return status in response.
    ///
    /// # Panics
    ///
    /// Panics if response json is corrupted.
    pub fn status(&mut self) -> Option<&i64> {
        self.result[0][0].into_i64()
    }

    /// Return start time in response.
    ///
    /// # Panics
    ///
    /// Panics if response json is corrupted.
    pub fn start_time(&mut self) -> Option<&f64> {
        self.result[0][1].into_f64()
    }

    /// Return elapsed time in response.
    ///
    /// # Panics
    ///
    /// Panics if response json is corrupted.
    pub fn elapsed_time(&mut self) -> Option<&f64> {
        self.result[0][2].into_f64()
    }

    /// [nodoc]
    #[inline]
    fn matched_columns_num(&mut self) -> Option<i64> {
        let vectoizable = match self.result[1][0][0].into_vec().clone() {
            Some(elem) => elem,
            None => return None
        };
        let pop = match vectoizable.clone().pop() {
            Some(pop) => pop,
            None => return None
        };
        match pop.into_i64().clone() {
            Some(v) => return Some(v.clone()),
            None => return None
        }
    }

    /// Return a number of matched columns in response.
    pub fn matched_columns(&mut self) -> Option<i64> {
        match self.status() {
            Some(&0) => self.matched_columns_num(),
            Some(_) => None,
            None    => None
        }
    }

    /// Get result in response.
    ///
    /// If request succeeded, it can get matched result array.
    /// Otherwise, one can get error messages.
    ///
    /// # Panics
    ///
    /// Panics if response json is corrupted.
    pub fn result(&mut self) -> Option<Vec<JFObject>> {
        match self.status() {
            Some(&0) => Some(vec![self.result[1][0].clone()]),
            Some(_) => Some(vec![self.result[0][3].clone()]),
            None    => None
        }
    }

    /// Convert to `Rows` type and return its type values.
    pub fn into_row(&mut self) -> Rows {
        Rows::new(self.result())
    }
}

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

    #[test]
    fn parse_result() {
        let response = "
    [[0,1452348610.39281,0.000101566314697266],
    [[[9],[[\"_id\",\"UInt32\"],[\"_key\",\"ShortText\"],[\"title\",\"ShortText\"]],
    [1,\"http://example.org/\",\"This is test record 1!\"],
    [2,\"http://example.net/\",\"test record 2.\"],
    [3,\"http://example.com/\",\"test test record three.\"],
    [4,\"http://example.net/afr\",\"test record four.\"],
    [5,\"http://example.org/aba\",\"test test test record five.\"],
    [6,\"http://example.com/rab\",\"test test test test record six.\"],
    [7,\"http://example.net/atv\",\"test test test record seven.\"],
    [8,\"http://example.org/gat\",\"test test record eight.\"],
    [9,\"http://example.com/vdw\",\"test test record nine.\"]]]]";
        let mut decode = ResultParser::new(response.to_string());
        assert_eq!(&0, decode.status().unwrap());
        assert_eq!(&1452348610.39281, decode.start_time().unwrap());
        assert_eq!(&0.000101566314697266, decode.elapsed_time().unwrap());
        assert_eq!(9, decode.matched_columns().unwrap());
        let vec = decode.result().unwrap().pop().unwrap().unwrap_vec().clone();
        let expected = r#"Array([Integer(1), String("http://example.org/"), String("This is test record 1!")])"#.to_owned();
        assert_eq!(expected, format!("{:?}", vec[2]))
    }
}
<commit_msg>rustdoc: Use link<commit_after>use json_flex;
use json_flex::JFObject;

#[derive(Clone)]
pub struct Rows {
    data: Option<Vec<JFObject>>
}

/// A result rows representing type.
/// This type is usually generated by [`ResultParser#into_raw()`](struct.ResultParser.html#method.into_row).
impl Rows {
    pub fn new(data: Option<Vec<JFObject>>) -> Rows {
        Rows { data: data }
    }

    /// Get columns in `Rows`.
    pub fn columns(&mut self) -> Option<Vec<JFObject>> {
        let popable = match self.data.clone() {
            Some(v) => v,
            None => return None
        };
        let pop = match popable.clone().pop() {
            Some(v) => v,
            None => return None
        };
        match pop.into_vec().clone() {
            Some(v) => return Some(v.clone()),
            None => return None
        }
    }
}

#[derive(Clone, Debug)]
pub struct ResultParser {
    result: Box<JFObject>
}

/// ResultParser
///
/// Groonga returns following array json:
///
/// success response:
/// `[[status,start_time, elapsed_time],
///   [[[matched_columns],[Array([column1,type1], ...)],
///     [Array([result1, result2, ...])]]]]`
///
/// error response:
/// `[[status,start_time, elapsed_time, error_information, ...]]`
impl ResultParser {
    pub fn new(json: String) -> ResultParser {
        ResultParser{result: json_flex::decode(json)}
    }

    /// Get raw response result.
    pub fn get_raw_object(&mut self) -> Box<JFObject> {
        self.result.clone()
    }

    /// Return header elements in response.
    ///
    /// # Panics
    ///
    /// Panics if response json is corrupted.
    pub fn get_header(&mut self) -> JFObject {
        self.result[0].clone()
    }

    /// Return status in response.
    ///
    /// # Panics
    ///
    /// Panics if response json is corrupted.
    pub fn status(&mut self) -> Option<&i64> {
        self.result[0][0].into_i64()
    }

    /// Return start time in response.
    ///
    /// # Panics
    ///
    /// Panics if response json is corrupted.
    pub fn start_time(&mut self) -> Option<&f64> {
        self.result[0][1].into_f64()
    }

    /// Return elapsed time in response.
    ///
    /// # Panics
    ///
    /// Panics if response json is corrupted.
    pub fn elapsed_time(&mut self) -> Option<&f64> {
        self.result[0][2].into_f64()
    }

    /// [nodoc]
    #[inline]
    fn matched_columns_num(&mut self) -> Option<i64> {
        let vectoizable = match self.result[1][0][0].into_vec().clone() {
            Some(elem) => elem,
            None => return None
        };
        let pop = match vectoizable.clone().pop() {
            Some(pop) => pop,
            None => return None
        };
        match pop.into_i64().clone() {
            Some(v) => return Some(v.clone()),
            None => return None
        }
    }

    /// Return a number of matched columns in response.
    pub fn matched_columns(&mut self) -> Option<i64> {
        match self.status() {
            Some(&0) => self.matched_columns_num(),
            Some(_) => None,
            None    => None
        }
    }

    /// Get result in response.
    ///
    /// If request succeeded, it can get matched result array.
    /// Otherwise, one can get error messages.
    ///
    /// # Panics
    ///
    /// Panics if response json is corrupted.
    pub fn result(&mut self) -> Option<Vec<JFObject>> {
        match self.status() {
            Some(&0) => Some(vec![self.result[1][0].clone()]),
            Some(_) => Some(vec![self.result[0][3].clone()]),
            None    => None
        }
    }

    /// Convert to `Rows` type and return its type values.
    pub fn into_row(&mut self) -> Rows {
        Rows::new(self.result())
    }
}

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

    #[test]
    fn parse_result() {
        let response = "
    [[0,1452348610.39281,0.000101566314697266],
    [[[9],[[\"_id\",\"UInt32\"],[\"_key\",\"ShortText\"],[\"title\",\"ShortText\"]],
    [1,\"http://example.org/\",\"This is test record 1!\"],
    [2,\"http://example.net/\",\"test record 2.\"],
    [3,\"http://example.com/\",\"test test record three.\"],
    [4,\"http://example.net/afr\",\"test record four.\"],
    [5,\"http://example.org/aba\",\"test test test record five.\"],
    [6,\"http://example.com/rab\",\"test test test test record six.\"],
    [7,\"http://example.net/atv\",\"test test test record seven.\"],
    [8,\"http://example.org/gat\",\"test test record eight.\"],
    [9,\"http://example.com/vdw\",\"test test record nine.\"]]]]";
        let mut decode = ResultParser::new(response.to_string());
        assert_eq!(&0, decode.status().unwrap());
        assert_eq!(&1452348610.39281, decode.start_time().unwrap());
        assert_eq!(&0.000101566314697266, decode.elapsed_time().unwrap());
        assert_eq!(9, decode.matched_columns().unwrap());
        let vec = decode.result().unwrap().pop().unwrap().unwrap_vec().clone();
        let expected = r#"Array([Integer(1), String("http://example.org/"), String("This is test record 1!")])"#.to_owned();
        assert_eq!(expected, format!("{:?}", vec[2]))
    }
}
<|endoftext|>


<commit_before>#![cfg(target_os = "linux")]

mod common;

use std::mem::size_of;
use std::thread;
use std::time::Duration;

use rand::Rng;

use sled::Config;

use common::cleanup;

const N_TESTS: usize = 100;
const CYCLE: usize = 256;
const BATCH_SIZE: u32 = 8;

/// Verifies that the keys in the tree are correctly recovered.
/// Panics if they are incorrect.
/// Returns the key that should be resumed at, and the current cycle value.
fn verify(tree: &sled::Tree) -> (u32, u32) {
    // key 0 should always be the highest value, as that's where we increment
    // at some point, it might go down by one
    // it should never return, or go down again after that
    let mut iter = tree.iter();
    let highest = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return (0, 0),
    };

    let highest_vec = u32_to_vec(highest);

    // find how far we got
    let mut contiguous: u32 = 0;
    let mut lowest = 0;
    for res in iter {
        let (_k, v) = res.unwrap();
        if &v[..4] == &highest_vec[..4] {
            contiguous += 1;
        } else {
            let expected = if highest == 0 {
                CYCLE as u32 - 1
            } else {
                (highest - 1) % CYCLE as u32
            };
            let actual = slice_to_u32(&*v);
            assert_eq!(expected, actual);
            lowest = actual;
            break;
        }
    }

    // ensure nothing changes after this point
    let low_beginning = u32_to_vec(contiguous + 1);

    for res in tree.range(&*low_beginning..) {
        let (k, v): (sled::IVec, _) = res.unwrap();
        assert_eq!(
            slice_to_u32(&*v),
            lowest,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            slice_to_u32(&*k),
            lowest,
            slice_to_u32(&*v),
            tree
        );
    }

    (contiguous, highest)
}

/// Verifies that the keys in the tree are correctly recovered (i.e., equal).
/// Panics if they are incorrect.
fn verify_batches(tree: &sled::Tree) -> u32 {
    let mut iter = tree.iter();
    let first_value = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return 0,
    };
    for key in 0..BATCH_SIZE {
        let res = tree.get(u32_to_vec(key));
        let option = res.unwrap();
        let v = match option {
            Some(v) => v,
            None => panic!(
                "expected key {} to have a value, instead it was missing in db: {:?}",
                key,
                tree
            ),
        };
        let value = slice_to_u32(&*v);
        assert_eq!(
            first_value,
            value,
            "expected key {} to have value {}, instead it had value {} in db: {:?}",
            key,
            first_value,
            value,
            tree
        );
    }

    first_value
}

fn u32_to_vec(u: u32) -> Vec<u8> {
    let buf: [u8; size_of::<u32>()] = u.to_be_bytes();
    buf.to_vec()
}

fn slice_to_u32(b: &[u8]) -> u32 {
    let mut buf = [0u8; size_of::<u32>()];
    buf.copy_from_slice(&b[..size_of::<u32>()]);

    u32::from_be_bytes(buf)
}

fn spawn_killah() {
    thread::spawn(|| {
        let runtime = rand::thread_rng().gen_range(0, 200);
        thread::sleep(Duration::from_millis(runtime));
        unsafe {
            libc::raise(9);
        }
    });
}

fn run(config: Config) {
    common::setup_logger();

    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = config.open().unwrap();

    if !crash_during_initialization {
        spawn_killah();
    }

    let (key, highest) = verify(&tree);

    let mut hu = ((highest as usize) * CYCLE) + key as usize;
    assert_eq!(hu % CYCLE, key as usize);
    assert_eq!(hu / CYCLE, highest as usize);

    loop {
        hu += 1;

        if hu / CYCLE >= CYCLE {
            hu = 0;
        }

        let key = u32_to_vec((hu % CYCLE) as u32);

        let mut value = u32_to_vec((hu / CYCLE) as u32);
        let additional_len = rand::thread_rng().gen_range(0, 1000);
        value.append(&mut vec![0u8; additional_len]);

        tree.insert(&key, value).unwrap();
    }
}

fn run_batches(config: Config) {
    fn do_batch(i: u32, tree: &sled::Tree) {
        let mut rng = rand::thread_rng();
        let base_value = u32_to_vec(i);

        let mut batch = sled::Batch::default();
        if rng.gen_bool(0.1) {
            for key in 0..BATCH_SIZE {
                batch.remove(u32_to_vec(key));
            }
        } else {
            for key in 0..BATCH_SIZE {
                let mut value = base_value.clone();
                let additional_len = rng.gen_range(0, 1000);
                value.append(&mut vec![0u8; additional_len]);

                batch.insert(u32_to_vec(key), value);
            }
        }
        tree.apply_batch(batch).unwrap();
    }

    common::setup_logger();
    let crash_during_initialization = rand::thread_rng().gen_bool(0.1);

    if crash_during_initialization {
        spawn_killah();
    }

    let tree = config.open().unwrap();

    let mut i = verify_batches(&tree);
    i += 1;
    do_batch(i, &tree);

    if !crash_during_initialization {
        spawn_killah();
    }

    loop {
        i += 1;
        do_batch(i, &tree);
    }
}

fn run_without_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56)
        .segment_size(1024);

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_with_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000)
        .segment_size(1024);

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_batches_without_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(1 << 56)
        .segment_size(1024);

    match thread::spawn(|| run_batches(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_batches_with_snapshot(dir: &str) {
    let config = Config::new()
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        .path(dir.to_string())
        .snapshot_after_ops(5000)
        .segment_size(1024);

    match thread::spawn(|| run_batches(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

#[test]
fn test_crash_recovery_with_runtime_snapshot() {
    let dir = "test_crashes_with_snapshot";
    cleanup(dir);
    for _ in 0..N_TESTS {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_with_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}

#[test]
fn test_crash_recovery_no_runtime_snapshot() {
    let dir = "test_crashes";
    cleanup(dir);
    for _ in 0..N_TESTS {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_without_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}

#[test]
#[ignore]
fn test_crash_batches_with_runtime_snapshot() {
    let dir = "test_batches_with_snapshot";
    cleanup(dir);
    for _ in 0..N_TESTS {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_batches_with_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}

#[test]
#[ignore]
fn test_crash_batches_no_runtime_snapshot() {
    let dir = "test_batches";
    cleanup(dir);
    for _ in 0..N_TESTS {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_batches_without_snapshot(dir)
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup(dir);
                panic!("child exited abnormally");
            }
        }
    }
    cleanup(dir);
}
<commit_msg>Remove old crash tests file<commit_after><|endoftext|>


<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use core_foundation::base::{CFRelease, CFRetain, CFTypeID, CFTypeRef, TCFType};

use libc::{c_void, size_t};
use std::mem;
use std::ptr;

pub type CGDataProviderGetBytesCallback = *const u8;
pub type CGDataProviderReleaseInfoCallback = *const u8;
pub type CGDataProviderRewindCallback = *const u8;
pub type CGDataProviderSkipBytesCallback = *const u8;
pub type CGDataProviderSkipForwardCallback = *const u8;

pub type CGDataProviderGetBytePointerCallback = *const u8;
pub type CGDataProviderGetBytesAtOffsetCallback = *const u8;
pub type CGDataProviderReleaseBytePointerCallback = *const u8;
pub type CGDataProviderReleaseDataCallback = *const u8;
pub type CGDataProviderGetBytesAtPositionCallback = *const u8;

#[repr(C)]
pub struct __CGDataProvider;

pub type CGDataProviderRef = *const __CGDataProvider;

pub struct CGDataProvider {
    obj: CGDataProviderRef,
}

impl Drop for CGDataProvider {
    fn drop(&mut self) {
        unsafe {
            CFRelease(self.as_CFTypeRef())
        }
    }
}

impl TCFType<CGDataProviderRef> for CGDataProvider {
    #[inline]
    fn as_concrete_TypeRef(&self) -> CGDataProviderRef {
        self.obj
    }

    #[inline]
    unsafe fn wrap_under_get_rule(reference: CGDataProviderRef) -> CGDataProvider {
        let reference: CGDataProviderRef = mem::transmute(CFRetain(mem::transmute(reference)));
        TCFType::wrap_under_create_rule(reference)
    }

    #[inline]
    fn as_CFTypeRef(&self) -> CFTypeRef {
        unsafe {
            mem::transmute(self.as_concrete_TypeRef())
        }
    }

    #[inline]
    unsafe fn wrap_under_create_rule(obj: CGDataProviderRef) -> CGDataProvider {
        CGDataProvider {
            obj: obj,
        }
    }

    #[inline]
    fn type_id() -> CFTypeID {
        unsafe {
            CGDataProviderGetTypeID()
        }
    }
}

impl CGDataProvider {
    pub fn from_buffer(buffer: &[u8]) -> CGDataProvider {
        unsafe {
            let result = CGDataProviderCreateWithData(ptr::null_mut(),
                                                      buffer.as_ptr() as *const c_void,
                                                      buffer.len() as size_t,
                                                      ptr::null());
            TCFType::wrap_under_create_rule(result)
        }
    }
}

#[link(name = "ApplicationServices", kind = "framework")]
extern {
    //fn CGDataProviderCopyData
    //fn CGDataProviderCreateDirect
    //fn CGDataProviderCreateSequential
    //fn CGDataProviderCreateWithCFData
    fn CGDataProviderCreateWithData(info: *mut c_void,
                                    data: *const c_void,
                                    size: size_t,
                                    releaseData: CGDataProviderReleaseDataCallback
                                   ) -> CGDataProviderRef;
    //fn CGDataProviderCreateWithFilename(filename: *c_char) -> CGDataProviderRef;
    //fn CGDataProviderCreateWithURL
    fn CGDataProviderGetTypeID() -> CFTypeID;
    //fn CGDataProviderRelease(provider: CGDataProviderRef);
    //fn CGDataProviderRetain(provider: CGDataProviderRef) -> CGDataProviderRef;
}
<commit_msg>Fix callback types<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use core_foundation::base::{CFRelease, CFRetain, CFTypeID, CFTypeRef, TCFType};

use libc::{c_void, size_t, off_t};
use std::mem;
use std::ptr;

pub type CGDataProviderGetBytesCallback = Option<unsafe extern fn (*mut c_void, *mut c_void, size_t) -> size_t>;
pub type CGDataProviderReleaseInfoCallback = Option<unsafe extern fn (*mut c_void)>;
pub type CGDataProviderRewindCallback = Option<unsafe extern fn (*mut c_void)>;
pub type CGDataProviderSkipBytesCallback = Option<unsafe extern fn (*mut c_void, size_t)>;
pub type CGDataProviderSkipForwardCallback = Option<unsafe extern fn (*mut c_void, off_t) -> off_t>;

pub type CGDataProviderGetBytePointerCallback = Option<unsafe extern fn (*mut c_void) -> *mut c_void>;
pub type CGDataProviderGetBytesAtOffsetCallback = Option<unsafe extern fn (*mut c_void, *mut c_void, size_t, size_t)>;
pub type CGDataProviderReleaseBytePointerCallback = Option<unsafe extern fn (*mut c_void, *const c_void)>;
pub type CGDataProviderReleaseDataCallback = Option<unsafe extern fn (*mut c_void, *const c_void, size_t)>;
pub type CGDataProviderGetBytesAtPositionCallback = Option<unsafe extern fn (*mut c_void, *mut c_void, off_t, size_t)>;

#[repr(C)]
pub struct __CGDataProvider;

pub type CGDataProviderRef = *const __CGDataProvider;

pub struct CGDataProvider {
    obj: CGDataProviderRef,
}

impl Drop for CGDataProvider {
    fn drop(&mut self) {
        unsafe {
            CFRelease(self.as_CFTypeRef())
        }
    }
}

impl TCFType<CGDataProviderRef> for CGDataProvider {
    #[inline]
    fn as_concrete_TypeRef(&self) -> CGDataProviderRef {
        self.obj
    }

    #[inline]
    unsafe fn wrap_under_get_rule(reference: CGDataProviderRef) -> CGDataProvider {
        let reference: CGDataProviderRef = mem::transmute(CFRetain(mem::transmute(reference)));
        TCFType::wrap_under_create_rule(reference)
    }

    #[inline]
    fn as_CFTypeRef(&self) -> CFTypeRef {
        unsafe {
            mem::transmute(self.as_concrete_TypeRef())
        }
    }

    #[inline]
    unsafe fn wrap_under_create_rule(obj: CGDataProviderRef) -> CGDataProvider {
        CGDataProvider {
            obj: obj,
        }
    }

    #[inline]
    fn type_id() -> CFTypeID {
        unsafe {
            CGDataProviderGetTypeID()
        }
    }
}

impl CGDataProvider {
    pub fn from_buffer(buffer: &[u8]) -> CGDataProvider {
        unsafe {
            let result = CGDataProviderCreateWithData(ptr::null_mut(),
                                                      buffer.as_ptr() as *const c_void,
                                                      buffer.len() as size_t,
                                                      None);
            TCFType::wrap_under_create_rule(result)
        }
    }
}

#[link(name = "ApplicationServices", kind = "framework")]
extern {
    //fn CGDataProviderCopyData
    //fn CGDataProviderCreateDirect
    //fn CGDataProviderCreateSequential
    //fn CGDataProviderCreateWithCFData
    fn CGDataProviderCreateWithData(info: *mut c_void,
                                    data: *const c_void,
                                    size: size_t,
                                    releaseData: CGDataProviderReleaseDataCallback
                                   ) -> CGDataProviderRef;
    //fn CGDataProviderCreateWithFilename(filename: *c_char) -> CGDataProviderRef;
    //fn CGDataProviderCreateWithURL
    fn CGDataProviderGetTypeID() -> CFTypeID;
    //fn CGDataProviderRelease(provider: CGDataProviderRef);
    //fn CGDataProviderRetain(provider: CGDataProviderRef) -> CGDataProviderRef;
}
<|endoftext|>


<commit_before>use std::io::{Error, ErrorKind};
use std::io;
use std::sync::mpsc::{self, Sender, Receiver};

use mio;
use mio::{EventLoop, EventSet, PollOpt, Handler, Token, TryWrite, TryRead};
use mio::tcp::{TcpListener, TcpStream};
use mio::util::Slab;
use bytes::{Buf, ByteBuf, MutByteBuf, SliceBuf};

use ::{SrvReq, SrvRes, CliReq, CliRes};
use codec;
use codec::Codec;

const SERVER_CLIENTS: Token = Token(0);
const SERVER_PEERS: Token = Token(1);

pub enum Message {
    PeerReq(SrvReq),
    PeerRes(SrvRes),
    CliReq(CliReq),
    CliRes(CliRes),
}

pub struct Envelope {
    id: u64,
    tok: Token,
    msg: ByteBuf,
}

pub struct Server {
    peers: Vec<String>,
    cli_handler: ConnSet,
    peer_handler: ConnSet,
}

impl Server {
    pub fn new(
        peer_port: u16,
        cli_port: u16,
        peers: Vec<String>,
        req_tx: Sender<Envelope>,
    ) -> io::Result<Server> {

        let cli_addr =
            format!("0.0.0.0:{}", cli_port).parse().unwrap();
        let cli_srv_sock =
            try!(TcpListener::bind(&cli_addr));

        let peer_addr =
            format!("0.0.0.0:{}", peer_port).parse().unwrap();
        let peer_srv_sock =
            try!(TcpListener::bind(&peer_addr));

        Ok(Server {
            peers: peers,
            cli_handler: ConnSet {
                srv_sock: cli_srv_sock,
                srv_token: SERVER_CLIENTS,
                conns: Slab::new_starting_at(Token(1024), 4096),
                req_tx: req_tx.clone(),
            },
            peer_handler: ConnSet {
                srv_sock: peer_srv_sock,
                srv_token: SERVER_PEERS,
                conns: Slab::new_starting_at(Token(2), 15),
                req_tx: req_tx.clone(),
            },
        })
    }

    pub fn run_event_loop(&mut self, mut event_loop: EventLoop<Server>) -> io::Result<()> {
        event_loop.register_opt(
            &self.cli_handler.srv_sock,
            SERVER_CLIENTS,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        ).unwrap();

        event_loop.register_opt(
            &self.peer_handler.srv_sock,
            SERVER_PEERS,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        ).unwrap();

        event_loop.run(self).unwrap();

        Err(Error::new(ErrorKind::Other, "event_loop shouldn't have returned."))

    }

    fn tok_to_sc(&mut self, tok: Token) -> Option<&mut ServerConn> {
        if tok.as_usize() > 1 && tok.as_usize() <= 128 {
            self.peer_handler.conns.get_mut(tok)
        } else if tok.as_usize() > 128 && tok.as_usize() <= 4096 {
            self.cli_handler.conns.get_mut(tok)
        } else {
            error!("bad event loop notification message envelope");
            None
        }
    }
}

impl Handler for Server {
    type Timeout = ();
    type Message = Envelope;

    fn ready(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        token: Token,
        events: EventSet,
    ) {
        if events.is_error() || events.is_hup() {
            println!("clearing error or hup connection");
            match token {
                peer if peer.as_usize() >= 2 && peer.as_usize() <= 16 => {
                    //self.peer_handler.conns.remove(token);
                },
                cli if cli.as_usize() >= 1024 && cli.as_usize() <= 4096 => {
                    //self.cli_handler.conns.remove(token);
                },
                t => panic!("bad token for error/hup: {}", t.as_usize()),
            }
        }

        if events.is_readable() {
            match token {
                SERVER_PEERS => {
                    info!("got SERVER_PEERS accept");
                    self.peer_handler.accept(event_loop).or_else( |e| {
                        error!("failed to accept peer: all slots full");
                        Err(e)
                    });
                },
                SERVER_CLIENTS => {
                    info!("got SERVER_CLIENTS accept");
                    self.cli_handler.accept(event_loop).or_else( |e| {
                        error!("failed to accept client: all slots full");
                        Err(e)
                    });
                },
                peer if peer.as_usize() >= 2 && peer.as_usize() <= 16 => {
                    self.peer_handler.conn_readable(event_loop, peer).unwrap();
                },
                cli if cli.as_usize() >= 1024 && cli.as_usize() <= 4096 => {
                    self.cli_handler.conn_readable(event_loop, cli).unwrap();
                },
                t => panic!("unknown token: {}", t.as_usize()),
            }
        }

        if events.is_writable() {
            match token {
                SERVER_PEERS =>
                    panic!("received writable for SERVER_PEERS"),
                SERVER_CLIENTS =>
                    panic!("received writable for token SERVER_CLIENTS"),
                peer if peer.as_usize() > 1 && peer.as_usize() <= 128 =>
                    self.peer_handler.conn_writable(event_loop, peer),
                cli if cli.as_usize() > 128 && cli.as_usize() <= 4096 =>
                    self.cli_handler.conn_writable(event_loop, cli),
                t => panic!("received writable for out-of-range token: {}", t.as_usize()),
            };
        }
    }

    fn notify(&mut self, event_loop: &mut EventLoop<Server>, mut msg: Envelope) {
        let mut sc = self.tok_to_sc(msg.tok).unwrap();
        // TODO(tyler) serialize <id> | <proto> and write to sc.res_buf
        let m = msg.msg.bytes();
        println!("msglen: {}", m.len());
        println!("bytes: {:?}", m);
        if sc.res_buf.is_none() {
            println!("1");

            let mut res = ByteBuf::mut_with_capacity(4 + m.len());
            assert!(res.write_slice(&codec::usize_to_array(m.len())) == 4);
            assert!(res.write_slice(m) == m.len());
            println!("sc.res_buf: {:?}", res.bytes());

            sc.res_remaining += res.bytes().len();
            sc.res_buf = Some(res.flip());
        } else {
            println!("2");
            // this shouldn't happen often, make a new one
            // TODO(tyler) concat bufs
            let mut rb = sc.res_buf.take().unwrap().flip();
            rb.write_slice(b"<id> | <proto response>\n");
            sc.res_buf = Some(rb.flip());
        }
        sc.interest.insert(EventSet::writable());
        event_loop.reregister(
            &sc.sock,
            msg.tok,
            sc.interest,
            PollOpt::edge() | PollOpt::oneshot(),
        );
    }
}

struct ServerConn {
    sock: TcpStream,
    req_tx: Sender<Envelope>,
    req_mut_buf: Option<MutByteBuf>,
    res_buf: Option<ByteBuf>,
    res_remaining: usize,
    req_codec: codec::Framed,
    token: Option<Token>,
    interest: EventSet
}

impl ServerConn {
    fn new(sock: TcpStream, req_tx: Sender<Envelope>) -> ServerConn {
        ServerConn {
            sock: sock,
            req_tx: req_tx,
            req_mut_buf: Some(ByteBuf::mut_with_capacity(2048)),
            req_codec: codec::Framed::new(),
            res_buf: None,
            res_remaining: 0,
            token: None,
            interest: EventSet::hup()
        }
    }

    fn writable(&mut self, event_loop: &mut EventLoop<Server>) -> io::Result<()> {
        if self.res_buf.is_none() {
            // no responses yet, don't reregister
            return Ok(())
        }
        let mut res_buf = self.res_buf.take().unwrap();

        println!("res buf: {:?}", res_buf.bytes());
        match self.sock.try_write_buf(&mut res_buf) {
            Ok(None) => {
                info!("client flushing buf; WOULDBLOCK");
                self.interest.insert(EventSet::writable());
            }
            Ok(Some(r)) => {
                info!("CONN : we wrote {} bytes!", r);
                println!("remaining: {}", self.res_remaining);
                self.res_remaining -= r;
                if self.res_remaining == 0 {
                    // we've written the whole response, now let's wait to read
                    self.interest.insert(EventSet::readable());
                    self.interest.remove(EventSet::writable());
                }
            }
            Err(e) => {
                match e.raw_os_error() {
                    Some(32) => {
                        info!("client disconnected");
                    },
                    Some(e) =>
                        info!("not implemented; client os err={:?}", e),
                    _ =>
                        info!("not implemented; client err={:?}", e),
                };
                // Don't reregister.
                return Err(e);
            },
        }
        self.interest.remove(EventSet::writable());

        self.res_buf = Some(res_buf);

        event_loop.reregister(
            &self.sock,
            self.token.unwrap(),
            self.interest,
            PollOpt::edge() | PollOpt::oneshot(),
        )
    }

    fn readable(&mut self, event_loop: &mut EventLoop<Server>) -> io::Result<()> {
        let mut req_buf = match self.req_mut_buf {
            Some(_) => self.req_mut_buf.take().unwrap(),
            None => ByteBuf::mut_with_capacity(2048),
        };

        match self.sock.try_read_buf(&mut req_buf) {
            Ok(None) => {
                panic!("got readable, but can't read from the socket");
            }
            Ok(Some(r)) => {
                info!("CONN : we read {} bytes!", r);
                self.interest.remove(EventSet::readable());
            }
            Err(e) => {
                info!("not implemented; client err={:?}", e);
                self.interest.remove(EventSet::readable());
            }
        };

        let mut rb = req_buf.flip();
        match self.req_codec.decode(&mut rb) {
            Some(req) => {
                self.req_tx.send(Envelope {
                    id: 5,
                    tok: self.token.unwrap(),
                    msg: req,
                });
            },
            None => {},
        }

        // TODO(tyler) this will get filled up quickly, move over remaining bytes after req
        self.req_mut_buf = Some(rb.flip());

        event_loop.reregister(
            &self.sock,
            self.token.unwrap(),
            self.interest,
            PollOpt::edge(),
        )
    }
}

pub struct ConnSet {
    srv_sock: TcpListener,
    srv_token: Token,
    conns: Slab<ServerConn>,
    req_tx: Sender<Envelope>,
}

impl ConnSet {
    fn accept(
        &mut self,
        event_loop: &mut EventLoop<Server>,
    ) -> io::Result<()> {

        info!("ConnSet accepting socket");

        let sock = self.srv_sock.accept().unwrap().unwrap();
        let conn = ServerConn::new(sock, self.req_tx.clone());

        // Re-register accepting socket
        event_loop.reregister(
            &self.srv_sock,
            self.srv_token,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        );

        self.conns.insert(conn).map(|tok| {
            // Register the connection
            self.conns[tok].token = Some(tok);
            event_loop.register_opt(
                &self.conns[tok].sock,
                tok,
                EventSet::readable(),
                PollOpt::edge() | PollOpt::oneshot()
            ).ok().expect("could not register socket with event loop");
            ()
        }).or_else(|e| Err(Error::new(ErrorKind::Other,
                                      "All connection slots full.")))
    }

    fn conn_readable(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        tok: Token,
    ) -> io::Result<()> {

        info!("ConnSet conn readable; tok={:?}", tok);
        self.conn(tok).readable(event_loop)
    }

    fn conn_writable(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        tok: Token,
    ) -> io::Result<()> {

        info!("ConnSet conn writable; tok={:?}", tok);
        match self.conn(tok).writable(event_loop) {
            Err(e) => {
                println!("got err in ConnSet conn_writable: {}", e);
                // now being done in server top level // self.conns.remove(tok);
                Err(e)
            },
            w => w,
        }
    }

    fn conn<'a>(&'a mut self, tok: Token) -> &'a mut ServerConn {
        &mut self.conns[tok]
    }
}
<commit_msg>fixed slab ServerConn leak<commit_after>use std::io::{Error, ErrorKind};
use std::io;
use std::sync::mpsc::{self, Sender, Receiver};

use mio;
use mio::{EventLoop, EventSet, PollOpt, Handler, Token, TryWrite, TryRead};
use mio::tcp::{TcpListener, TcpStream};
use mio::util::Slab;
use bytes::{Buf, ByteBuf, MutByteBuf, SliceBuf};

use ::{SrvReq, SrvRes, CliReq, CliRes};
use codec;
use codec::Codec;

const SERVER_CLIENTS: Token = Token(0);
const SERVER_PEERS: Token = Token(1);

pub enum Message {
    PeerReq(SrvReq),
    PeerRes(SrvRes),
    CliReq(CliReq),
    CliRes(CliRes),
}

pub struct Envelope {
    id: u64,
    tok: Token,
    msg: ByteBuf,
}

pub struct Server {
    peers: Vec<String>,
    cli_handler: ConnSet,
    peer_handler: ConnSet,
}

impl Server {
    pub fn new(
        peer_port: u16,
        cli_port: u16,
        peers: Vec<String>,
        req_tx: Sender<Envelope>,
    ) -> io::Result<Server> {

        let cli_addr =
            format!("0.0.0.0:{}", cli_port).parse().unwrap();
        let cli_srv_sock =
            try!(TcpListener::bind(&cli_addr));

        let peer_addr =
            format!("0.0.0.0:{}", peer_port).parse().unwrap();
        let peer_srv_sock =
            try!(TcpListener::bind(&peer_addr));

        Ok(Server {
            peers: peers,
            cli_handler: ConnSet {
                srv_sock: cli_srv_sock,
                srv_token: SERVER_CLIENTS,
                conns: Slab::new_starting_at(Token(1024), 4096),
                req_tx: req_tx.clone(),
            },
            peer_handler: ConnSet {
                srv_sock: peer_srv_sock,
                srv_token: SERVER_PEERS,
                conns: Slab::new_starting_at(Token(2), 15),
                req_tx: req_tx.clone(),
            },
        })
    }

    pub fn run_event_loop(&mut self, mut event_loop: EventLoop<Server>) -> io::Result<()> {
        event_loop.register_opt(
            &self.cli_handler.srv_sock,
            SERVER_CLIENTS,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        ).unwrap();

        event_loop.register_opt(
            &self.peer_handler.srv_sock,
            SERVER_PEERS,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        ).unwrap();

        event_loop.run(self).unwrap();

        Err(Error::new(ErrorKind::Other, "event_loop shouldn't have returned."))

    }

    fn tok_to_sc(&mut self, tok: Token) -> Option<&mut ServerConn> {
        if tok.as_usize() > 1 && tok.as_usize() <= 128 {
            self.peer_handler.conns.get_mut(tok)
        } else if tok.as_usize() > 128 && tok.as_usize() <= 4096 {
            self.cli_handler.conns.get_mut(tok)
        } else {
            error!("bad event loop notification message envelope");
            None
        }
    }
}

impl Handler for Server {
    type Timeout = ();
    type Message = Envelope;

    fn ready(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        token: Token,
        events: EventSet,
    ) {
        if events.is_error() || events.is_hup() {
            println!("clearing error or hup connection");
            match token {
                peer if peer.as_usize() >= 2 && peer.as_usize() <= 16 => {
                    if self.peer_handler.conns.contains(token) {
                         self.peer_handler.conns.remove(token);
                    }
                },
                cli if cli.as_usize() >= 1024 && cli.as_usize() <= 4096 => {
                    if self.cli_handler.conns.contains(token) {
                        self.cli_handler.conns.remove(token);
                    }
                },
                t => panic!("bad token for error/hup: {}", t.as_usize()),
            }
        }

        if events.is_readable() {
            match token {
                SERVER_PEERS => {
                    info!("got SERVER_PEERS accept");
                    self.peer_handler.accept(event_loop).or_else( |e| {
                        error!("failed to accept peer: all slots full");
                        Err(e)
                    });
                },
                SERVER_CLIENTS => {
                    info!("got SERVER_CLIENTS accept");
                    self.cli_handler.accept(event_loop).or_else( |e| {
                        error!("failed to accept client: all slots full");
                        Err(e)
                    });
                },
                peer if peer.as_usize() >= 2 && peer.as_usize() <= 16 => {
                    self.peer_handler.conn_readable(event_loop, peer).unwrap();
                },
                cli if cli.as_usize() >= 1024 && cli.as_usize() <= 4096 => {
                    self.cli_handler.conn_readable(event_loop, cli).unwrap();
                },
                t => panic!("unknown token: {}", t.as_usize()),
            }
        }

        if events.is_writable() {
            match token {
                SERVER_PEERS =>
                    panic!("received writable for SERVER_PEERS"),
                SERVER_CLIENTS =>
                    panic!("received writable for token SERVER_CLIENTS"),
                peer if peer.as_usize() > 1 && peer.as_usize() <= 128 =>
                    self.peer_handler.conn_writable(event_loop, peer),
                cli if cli.as_usize() > 128 && cli.as_usize() <= 4096 =>
                    self.cli_handler.conn_writable(event_loop, cli),
                t => panic!("received writable for out-of-range token: {}", t.as_usize()),
            };
        }
    }

    fn notify(&mut self, event_loop: &mut EventLoop<Server>, mut msg: Envelope) {
        let mut sc = self.tok_to_sc(msg.tok).unwrap();
        // TODO(tyler) serialize <id> | <proto> and write to sc.res_buf
        let m = msg.msg.bytes();
        println!("msglen: {}", m.len());
        println!("bytes: {:?}", m);
        if sc.res_buf.is_none() {
            println!("1");

            let mut res = ByteBuf::mut_with_capacity(4 + m.len());
            assert!(res.write_slice(&codec::usize_to_array(m.len())) == 4);
            assert!(res.write_slice(m) == m.len());
            println!("sc.res_buf: {:?}", res.bytes());

            sc.res_remaining += res.bytes().len();
            sc.res_buf = Some(res.flip());
        } else {
            println!("2");
            // this shouldn't happen often, make a new one
            // TODO(tyler) concat bufs
            let mut rb = sc.res_buf.take().unwrap().flip();
            rb.write_slice(b"<id> | <proto response>\n");
            sc.res_buf = Some(rb.flip());
        }
        sc.interest.insert(EventSet::writable());
        event_loop.reregister(
            &sc.sock,
            msg.tok,
            sc.interest,
            PollOpt::edge() | PollOpt::oneshot(),
        );
    }
}

struct ServerConn {
    sock: TcpStream,
    req_tx: Sender<Envelope>,
    req_mut_buf: Option<MutByteBuf>,
    res_buf: Option<ByteBuf>,
    res_remaining: usize,
    req_codec: codec::Framed,
    token: Option<Token>,
    interest: EventSet
}

impl ServerConn {
    fn new(sock: TcpStream, req_tx: Sender<Envelope>) -> ServerConn {
        ServerConn {
            sock: sock,
            req_tx: req_tx,
            req_mut_buf: Some(ByteBuf::mut_with_capacity(2048)),
            req_codec: codec::Framed::new(),
            res_buf: None,
            res_remaining: 0,
            token: None,
            interest: EventSet::hup()
        }
    }

    fn writable(&mut self, event_loop: &mut EventLoop<Server>) -> io::Result<()> {
        if self.res_buf.is_none() {
            // no responses yet, don't reregister
            return Ok(())
        }
        let mut res_buf = self.res_buf.take().unwrap();

        println!("res buf: {:?}", res_buf.bytes());
        match self.sock.try_write_buf(&mut res_buf) {
            Ok(None) => {
                info!("client flushing buf; WOULDBLOCK");
                self.interest.insert(EventSet::writable());
            }
            Ok(Some(r)) => {
                info!("CONN : we wrote {} bytes!", r);
                println!("remaining: {}", self.res_remaining);
                self.res_remaining -= r;
                if self.res_remaining == 0 {
                    // we've written the whole response, now let's wait to read
                    self.interest.insert(EventSet::readable());
                    self.interest.remove(EventSet::writable());
                }
            }
            Err(e) => {
                match e.raw_os_error() {
                    Some(32) => {
                        info!("client disconnected");
                    },
                    Some(e) =>
                        info!("not implemented; client os err={:?}", e),
                    _ =>
                        info!("not implemented; client err={:?}", e),
                };
                // Don't reregister.
                return Err(e);
            },
        }
        self.interest.remove(EventSet::writable());

        self.res_buf = Some(res_buf);

        event_loop.reregister(
            &self.sock,
            self.token.unwrap(),
            self.interest,
            PollOpt::edge() | PollOpt::oneshot(),
        )
    }

    fn readable(&mut self, event_loop: &mut EventLoop<Server>) -> io::Result<()> {
        let mut req_buf = match self.req_mut_buf {
            Some(_) => self.req_mut_buf.take().unwrap(),
            None => ByteBuf::mut_with_capacity(2048),
        };

        match self.sock.try_read_buf(&mut req_buf) {
            Ok(None) => {
                panic!("got readable, but can't read from the socket");
            }
            Ok(Some(r)) => {
                info!("CONN : we read {} bytes!", r);
                self.interest.remove(EventSet::readable());
            }
            Err(e) => {
                info!("not implemented; client err={:?}", e);
                self.interest.remove(EventSet::readable());
            }
        };

        let mut rb = req_buf.flip();
        match self.req_codec.decode(&mut rb) {
            Some(req) => {
                self.req_tx.send(Envelope {
                    id: 5,
                    tok: self.token.unwrap(),
                    msg: req,
                });
            },
            None => {},
        }

        // TODO(tyler) this will get filled up quickly, move over remaining bytes after req
        self.req_mut_buf = Some(rb.flip());

        event_loop.reregister(
            &self.sock,
            self.token.unwrap(),
            self.interest,
            PollOpt::edge(),
        )
    }
}

pub struct ConnSet {
    srv_sock: TcpListener,
    srv_token: Token,
    conns: Slab<ServerConn>,
    req_tx: Sender<Envelope>,
}

impl ConnSet {
    fn accept(
        &mut self,
        event_loop: &mut EventLoop<Server>,
    ) -> io::Result<()> {

        info!("ConnSet accepting socket");

        let sock = self.srv_sock.accept().unwrap().unwrap();
        let conn = ServerConn::new(sock, self.req_tx.clone());

        // Re-register accepting socket
        event_loop.reregister(
            &self.srv_sock,
            self.srv_token,
            EventSet::readable(),
            PollOpt::edge() | PollOpt::oneshot(),
        );

        self.conns.insert(conn).map(|tok| {
            // Register the connection
            self.conns[tok].token = Some(tok);
            event_loop.register_opt(
                &self.conns[tok].sock,
                tok,
                EventSet::readable(),
                PollOpt::edge() | PollOpt::oneshot()
            ).ok().expect("could not register socket with event loop");
            ()
        }).or_else(|e| Err(Error::new(ErrorKind::Other,
                                      "All connection slots full.")))
    }

    fn conn_readable(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        tok: Token,
    ) -> io::Result<()> {

        info!("ConnSet conn readable; tok={:?}", tok);
        if !self.conns.contains(tok) {
            error!("got conn_readable for non-existent token!");
            return Ok(());
        }

        self.conn(tok).readable(event_loop)
    }

    fn conn_writable(
        &mut self,
        event_loop: &mut EventLoop<Server>,
        tok: Token,
    ) -> io::Result<()> {
        if !self.conns.contains(tok) {
            error!("got conn_readable for non-existent token!");
            return Ok(());
        }

        info!("ConnSet conn writable; tok={:?}", tok);
        match self.conn(tok).writable(event_loop) {
            Err(e) => {
                println!("got err in ConnSet conn_writable: {}", e);
                // now being done in server top level // self.conns.remove(tok);
                Err(e)
            },
            w => w,
        }
    }

    fn conn<'a>(&'a mut self, tok: Token) -> &'a mut ServerConn {
        &mut self.conns[tok]
    }
}
<|endoftext|>


<commit_before>use std::io;
use futures::{Future, Stream};
use futures::future::{err, ok};
use serde_json;
use model::{RegisterRequest, Instance};
use errors::EurekaClientError;
use hyper::{Client, Method, Request, Body, Uri, mime, Error as HyperError, StatusCode};
use hyper::header::{Accept, AcceptEncoding, Encoding, Headers, UserAgent, ContentType, ContentLength, qitem};
use tokio_core::reactor::Handle;

pub struct EurekaClient<'a> {
    handle: &'a Handle,
    client_name: String,
    eureka_cluster_url: String,
}

//
// A simple port of the example found at: https://github.com/Netflix/eureka/wiki/Example-Custom-ReadOnly-client
// Eureka REST API: https://github.com/Netflix/eureka/wiki/Eureka-REST-operations
impl<'a> EurekaClient<'a> {
    pub fn new(handle: &'a Handle, client_name: &str, eureka_cluster_url: &str) -> EurekaClient<'a> {
        debug!("Creating new Eureka Client client_name:{:?}, eureka_client:{:?}", client_name, eureka_cluster_url);
        EurekaClient {
            handle: &handle,
            client_name: client_name.to_owned(),
            eureka_cluster_url: eureka_cluster_url.to_owned()
        }
    }

    pub fn register(&self, application_id: &str, register_request: &RegisterRequest) -> Box<Future<Item=(), Error=EurekaClientError>> {
        debug!("register: application_id={:?}, register_request:{:?}", application_id, register_request);
        let client = Client::new(self.handle);
        let path = "/v2/apps/".to_owned() + application_id;
        let mut req: Request<Body> = Request::new(Method::Post, self.build_uri(path.as_ref()));
        self.set_headers(req.headers_mut());

        let json = serde_json::to_string(register_request).unwrap();
        req.headers_mut().set(ContentLength(json.len() as u64));
        req.set_body(json);

        let result = client.request(req)
            .map_err(|e| {
                EurekaClientError::from(e)
            })
            .and_then(|res| {
                debug!("register: server response {:?}", res);

                let status = res.status();
                match status {
                    StatusCode::BadRequest => Err(EurekaClientError::BadRequest),
                    StatusCode::InternalServerError => Err(EurekaClientError::InternalServerError),
                    _ => Ok(())
                }
            });
        Box::new(result)
    }

    pub fn get_application_instances<'b>(&self, application_id: &str) -> Box<Future<Item=Vec<Instance>, Error=EurekaClientError>> {
        let client = Client::new(self.handle);
        let path = "/v2/apps/".to_owned() + application_id;
        let mut req: Request<Body> = Request::new(Method::Get, self.build_uri(path.as_ref()));
        self.set_headers(req.headers_mut());

        let result = client.request(req).and_then(|res| {
            let status = res.status();
            debug!("get_application_instance: server response status: {}", status);
            res.body().concat2().and_then(move |body| {
                match status {
                    StatusCode::NotFound => Ok(vec![]),
                    _ => {
                        serde_json::from_slice::<Vec<Instance>>(&body).map_err(|e| {
                            HyperError::Io(io::Error::new(io::ErrorKind::Other, e))
                        })
                    }
                }
            })
        }).map_err(|e| {
            EurekaClientError::from(e)
        });
        Box::new(result)
    }

    fn build_uri(&self, path: &str) -> Uri {
        (self.eureka_cluster_url.to_owned() + path).parse().unwrap()
    }

    fn set_headers(&self, headers: &mut Headers) {
        headers.set(AcceptEncoding(vec![qitem(Encoding::Gzip)]));
        headers.set(Accept(vec![qitem(mime::APPLICATION_JSON)]));
        headers.set(ContentType(mime::APPLICATION_JSON));
        let user_agent = "Rust Hyper/".to_string() + self.client_name.as_ref();
        headers.set(UserAgent::new(user_agent));
    }
}
<commit_msg>figured out how to coerce my NotFoundError<commit_after>use std::io;
use futures::{Future, Stream};
use futures::future::{err, ok};
use serde_json;
use model::{RegisterRequest, Instance};
use errors::EurekaClientError;
use hyper::{Client, Method, Request, Body, Uri, mime, Error as HyperError, StatusCode};
use hyper::header::{Accept, AcceptEncoding, Encoding, Headers, UserAgent, ContentType, ContentLength, qitem};
use tokio_core::reactor::Handle;

pub struct EurekaClient<'a> {
    handle: &'a Handle,
    client_name: String,
    eureka_cluster_url: String,
}

//
// A simple port of the example found at: https://github.com/Netflix/eureka/wiki/Example-Custom-ReadOnly-client
// Eureka REST API: https://github.com/Netflix/eureka/wiki/Eureka-REST-operations
impl<'a> EurekaClient<'a> {
    pub fn new(handle: &'a Handle, client_name: &str, eureka_cluster_url: &str) -> EurekaClient<'a> {
        debug!("Creating new Eureka Client client_name:{:?}, eureka_client:{:?}", client_name, eureka_cluster_url);
        EurekaClient {
            handle: &handle,
            client_name: client_name.to_owned(),
            eureka_cluster_url: eureka_cluster_url.to_owned()
        }
    }

    pub fn register(&self, application_id: &str, register_request: &RegisterRequest) -> Box<Future<Item=(), Error=EurekaClientError>> {
        debug!("register: application_id={:?}, register_request:{:?}", application_id, register_request);
        let client = Client::new(self.handle);
        let path = "/v2/apps/".to_owned() + application_id;
        let mut req: Request<Body> = Request::new(Method::Post, self.build_uri(path.as_ref()));
        self.set_headers(req.headers_mut());

        let json = serde_json::to_string(register_request).unwrap();
        req.headers_mut().set(ContentLength(json.len() as u64));
        req.set_body(json);

        let result = client.request(req)
            .map_err(|e| {
                EurekaClientError::from(e)
            })
            .and_then(|res| {
                debug!("register: server response {:?}", res);

                let status = res.status();
                match status {
                    StatusCode::BadRequest => Err(EurekaClientError::BadRequest),
                    StatusCode::InternalServerError => Err(EurekaClientError::InternalServerError),
                    _ => Ok(())
                }
            });
        Box::new(result)
    }

    pub fn get_application_instances<'b>(&self, application_id: &str) -> Box<Future<Item=Vec<Instance>, Error=EurekaClientError>> {

        // Since it was hard to coerce the errot type into a EurekaClientError
        // I set the result in a holder then map result into an error or ok
        // There has to be a better way.. but this works.
        enum IntermediateResult {
            Ok(Vec<Instance>),
            Err(EurekaClientError)
        }
        impl IntermediateResult {
            fn ok(vec: Vec<Instance>) -> IntermediateResult {
                IntermediateResult::Ok(vec)
            }
            fn err(err: EurekaClientError) -> IntermediateResult {
                IntermediateResult::Err(err)
            }
        }

        let client = Client::new(self.handle);
        let path = "/v2/apps/".to_owned() + application_id;
        let mut req: Request<Body> = Request::new(Method::Get, self.build_uri(path.as_ref()));
        self.set_headers(req.headers_mut());

        let result = client.request(req).and_then(|res| {
            let status = res.status();
            debug!("get_application_instance: server response status: {}", status);
            res.body().concat2().and_then(move |body| {
                match status {
                    StatusCode::NotFound => Ok(IntermediateResult::err(EurekaClientError::NotFound)),
                    _ => {
                        serde_json::from_slice::<Vec<Instance>>(&body).map_err(|e| {
                            HyperError::Io(io::Error::new(io::ErrorKind::Other, e))
                        })
                        .map(|r| IntermediateResult::ok(r))
                    }
                }
            })
        })
        .map_err(|e| {
            EurekaClientError::from(e)
        })
        .and_then(|ir| {
            // now that we have changed the error to EurekaClientError
            // we can map our err back in
            match ir {
                IntermediateResult::Ok(vec) => Ok(vec),
                IntermediateResult::Err(err) => Err(err)
            }
        });
        Box::new(result)
    }

    fn build_uri(&self, path: &str) -> Uri {
        (self.eureka_cluster_url.to_owned() + path).parse().unwrap()
    }

    fn set_headers(&self, headers: &mut Headers) {
        headers.set(AcceptEncoding(vec![qitem(Encoding::Gzip)]));
        headers.set(Accept(vec![qitem(mime::APPLICATION_JSON)]));
        headers.set(ContentType(mime::APPLICATION_JSON));
        let user_agent = "Rust Hyper/".to_string() + self.client_name.as_ref();
        headers.set(UserAgent::new(user_agent));
    }
}
<|endoftext|>