<commit_before>use apply_rewards::process_rewards_and_penalties;
use errors::EpochProcessingError as Error;
use process_slashings::process_slashings;
use registry_updates::process_registry_updates;
use std::collections::HashMap;
use tree_hash::TreeHash;
use types::*;
use validator_statuses::{TotalBalances, ValidatorStatuses};
use winning_root::{winning_root, WinningRoot};

pub mod apply_rewards;
pub mod errors;
pub mod get_attesting_indices;
pub mod inclusion_distance;
pub mod process_slashings;
pub mod registry_updates;
pub mod tests;
pub mod validator_statuses;
pub mod winning_root;

/// Maps a shard to a winning root.
///
/// It is generated during crosslink processing and later used to reward/penalize validators.
pub type WinningRootHashSet = HashMap<u64, WinningRoot>;

/// Performs per-epoch processing on some BeaconState.
///
/// Mutates the given `BeaconState`, returning early if an error is encountered. If an error is
/// returned, a state might be "half-processed" and therefore in an invalid state.
///
/// Spec v0.6.1
pub fn per_epoch_processing<T: EthSpec>(
    state: &mut BeaconState<T>,
    spec: &ChainSpec,
) -> Result<(), Error> {
    // Ensure the previous and next epoch caches are built.
    state.build_epoch_cache(RelativeEpoch::Previous, spec)?;
    state.build_epoch_cache(RelativeEpoch::Current, spec)?;

    // Load the struct we use to assign validators into sets based on their participation.
    //
    // E.g., attestation in the previous epoch, attested to the head, etc.
    let mut validator_statuses = ValidatorStatuses::new(state, spec)?;
    validator_statuses.process_attestations(&state, spec)?;

    // Justification and finalization.
    process_justification_and_finalization(state, &validator_statuses.total_balances, spec)?;

    // Crosslinks.
    let winning_root_for_shards = process_crosslinks(state, spec)?;

    // Rewards and Penalities.
    process_rewards_and_penalties(
        state,
        &mut validator_statuses,
        &winning_root_for_shards,
        spec,
    )?;

    // Registry Updates.
    process_registry_updates(state, spec)?;

    // Slashings.
    process_slashings(state, validator_statuses.total_balances.current_epoch, spec)?;

    // Final updates.
    process_final_updates(state, spec)?;

    // Rotate the epoch caches to suit the epoch transition.
    state.advance_caches();

    Ok(())
}

/// Update the following fields on the `BeaconState`:
///
/// - `justification_bitfield`.
/// - `previous_justified_epoch`
/// - `previous_justified_root`
/// - `current_justified_epoch`
/// - `current_justified_root`
/// - `finalized_epoch`
/// - `finalized_root`
///
/// Spec v0.6.1
pub fn process_justification_and_finalization<T: EthSpec>(
    state: &mut BeaconState<T>,
    total_balances: &TotalBalances,
    spec: &ChainSpec,
) -> Result<(), Error> {
    if state.current_epoch(spec) == spec.genesis_epoch {
        return Ok(());
    }

    let previous_epoch = state.previous_epoch(spec);
    let current_epoch = state.current_epoch(spec);

    let old_previous_justified_epoch = state.previous_justified_epoch;
    let old_current_justified_epoch = state.current_justified_epoch;

    // Process justifications
    state.previous_justified_epoch = state.current_justified_epoch;
    state.previous_justified_root = state.current_justified_root;
    state.justification_bitfield <<= 1;

    if total_balances.previous_epoch_target_attesters * 3 >= total_balances.previous_epoch * 2 {
        state.current_justified_epoch = previous_epoch;
        state.current_justified_root =
            *state.get_block_root_at_epoch(state.current_justified_epoch, spec)?;
        state.justification_bitfield |= 2;
    }
    // If the current epoch gets justified, fill the last bit.
    if total_balances.current_epoch_target_attesters * 3 >= total_balances.current_epoch * 2 {
        state.current_justified_epoch = current_epoch;
        state.current_justified_root =
            *state.get_block_root_at_epoch(state.current_justified_epoch, spec)?;
        state.justification_bitfield |= 1;
    }

    let bitfield = state.justification_bitfield;

    // The 2nd/3rd/4th most recent epochs are all justified, the 2nd using the 4th as source.
    if (bitfield >> 1) % 8 == 0b111 && old_previous_justified_epoch == current_epoch - 3 {
        state.finalized_epoch = old_previous_justified_epoch;
        state.finalized_root = *state.get_block_root_at_epoch(state.finalized_epoch, spec)?;
    }
    // The 2nd/3rd most recent epochs are both justified, the 2nd using the 3rd as source.
    if (bitfield >> 1) % 4 == 0b11 && state.previous_justified_epoch == current_epoch - 2 {
        state.finalized_epoch = old_previous_justified_epoch;
        state.finalized_root = *state.get_block_root_at_epoch(state.finalized_epoch, spec)?;
    }
    // The 1st/2nd/3rd most recent epochs are all justified, the 1st using the 2nd as source.
    if bitfield % 8 == 0b111 && state.current_justified_epoch == current_epoch - 2 {
        state.finalized_epoch = old_current_justified_epoch;
        state.finalized_root = *state.get_block_root_at_epoch(state.finalized_epoch, spec)?;
    }
    // The 1st/2nd most recent epochs are both justified, the 1st using the 2nd as source.
    if bitfield % 4 == 0b11 && state.current_justified_epoch == current_epoch - 1 {
        state.finalized_epoch = old_current_justified_epoch;
        state.finalized_root = *state.get_block_root_at_epoch(state.finalized_epoch, spec)?;
    }

    Ok(())
}

/// Updates the following fields on the `BeaconState`:
///
/// - `previous_crosslinks`
/// - `current_crosslinks`
///
/// Also returns a `WinningRootHashSet` for later use during epoch processing.
///
/// Spec v0.6.1
pub fn process_crosslinks<T: EthSpec>(
    state: &mut BeaconState<T>,
    spec: &ChainSpec,
) -> Result<WinningRootHashSet, Error> {
    let mut winning_root_for_shards: WinningRootHashSet = HashMap::new();

    state.previous_crosslinks = state.current_crosslinks.clone();

    for epoch in vec![state.previous_epoch(spec), state.current_epoch(spec)] {
        for offset in 0..state.get_epoch_committee_count(epoch, spec) {
            let shard = (state.get_epoch_start_shard(epoch, spec) + offset) % spec.shard_count;
            let crosslink_committee = state.get_crosslink_committee(epoch, shard, spec)?;

            let winning_root = winning_root(state, shard, epoch, spec)?;

            if let Some(winning_root) = winning_root {
                let total_committee_balance =
                    state.get_total_balance(&crosslink_committee.committee, spec)?;

                if 3 * winning_root.total_attesting_balance >= 2 * total_committee_balance {
                    state.current_crosslinks[shard as usize] = winning_root.crosslink.clone();
                }
                winning_root_for_shards.insert(shard, winning_root);
            }
        }
    }

    Ok(winning_root_for_shards)
}

/// Finish up an epoch update.
///
/// Spec v0.6.1
pub fn process_final_updates<T: EthSpec>(
    state: &mut BeaconState<T>,
    spec: &ChainSpec,
) -> Result<(), Error> {
    let current_epoch = state.current_epoch(spec);
    let next_epoch = state.next_epoch(spec);

    // Reset eth1 data votes.
    if (state.slot + 1) % spec.slots_per_eth1_voting_period == 0 {
        state.eth1_data_votes = vec![];
    }

    // Update effective balances with hysteresis (lag).
    for (index, validator) in state.validator_registry.iter_mut().enumerate() {
        let balance = state.balances[index];
        let half_increment = spec.effective_balance_increment / 2;
        if balance < validator.effective_balance
            || validator.effective_balance + 3 * half_increment < balance
        {
            validator.effective_balance = std::cmp::min(
                balance - balance % spec.effective_balance_increment,
                spec.max_effective_balance,
            );
        }
    }

    // Update start shard.
    state.latest_start_shard = (state.latest_start_shard
        + state.get_shard_delta(current_epoch, spec))
        % T::ShardCount::to_u64();

    // This is a hack to allow us to update index roots and slashed balances for the next epoch.
    //
    // The indentation here is to make it obvious where the weird stuff happens.
    {
        state.slot += 1;

        // Set active index root
        let active_index_root = Hash256::from_slice(
            &state
                .get_active_validator_indices(next_epoch + spec.activation_exit_delay)
                .tree_hash_root()[..],
        );
        state.set_active_index_root(next_epoch, active_index_root, spec)?;

        // Set total slashed balances
        state.set_slashed_balance(next_epoch, state.get_slashed_balance(current_epoch)?)?;

        // Set randao mix
        state.set_randao_mix(
            next_epoch,
            *state.get_randao_mix(current_epoch, spec)?,
            spec,
        )?;

        state.slot -= 1;
    }

    if next_epoch.as_u64() % (T::SlotsPerHistoricalRoot::to_u64() / spec.slots_per_epoch) == 0 {
        let historical_batch = state.historical_batch();
        state
            .historical_roots
            .push(Hash256::from_slice(&historical_batch.tree_hash_root()[..]));
    }

    // Rotate current/previous epoch attestations
    std::mem::swap(
        &mut state.previous_epoch_attestations,
        &mut state.current_epoch_attestations,
    );
    state.current_epoch_attestations = vec![];

    Ok(())
}
<commit_msg>state_processing: nicer attestation swap<commit_after>use apply_rewards::process_rewards_and_penalties;
use errors::EpochProcessingError as Error;
use process_slashings::process_slashings;
use registry_updates::process_registry_updates;
use std::collections::HashMap;
use tree_hash::TreeHash;
use types::*;
use validator_statuses::{TotalBalances, ValidatorStatuses};
use winning_root::{winning_root, WinningRoot};

pub mod apply_rewards;
pub mod errors;
pub mod get_attesting_indices;
pub mod inclusion_distance;
pub mod process_slashings;
pub mod registry_updates;
pub mod tests;
pub mod validator_statuses;
pub mod winning_root;

/// Maps a shard to a winning root.
///
/// It is generated during crosslink processing and later used to reward/penalize validators.
pub type WinningRootHashSet = HashMap<u64, WinningRoot>;

/// Performs per-epoch processing on some BeaconState.
///
/// Mutates the given `BeaconState`, returning early if an error is encountered. If an error is
/// returned, a state might be "half-processed" and therefore in an invalid state.
///
/// Spec v0.6.1
pub fn per_epoch_processing<T: EthSpec>(
    state: &mut BeaconState<T>,
    spec: &ChainSpec,
) -> Result<(), Error> {
    // Ensure the previous and next epoch caches are built.
    state.build_epoch_cache(RelativeEpoch::Previous, spec)?;
    state.build_epoch_cache(RelativeEpoch::Current, spec)?;

    // Load the struct we use to assign validators into sets based on their participation.
    //
    // E.g., attestation in the previous epoch, attested to the head, etc.
    let mut validator_statuses = ValidatorStatuses::new(state, spec)?;
    validator_statuses.process_attestations(&state, spec)?;

    // Justification and finalization.
    process_justification_and_finalization(state, &validator_statuses.total_balances, spec)?;

    // Crosslinks.
    let winning_root_for_shards = process_crosslinks(state, spec)?;

    // Rewards and Penalities.
    process_rewards_and_penalties(
        state,
        &mut validator_statuses,
        &winning_root_for_shards,
        spec,
    )?;

    // Registry Updates.
    process_registry_updates(state, spec)?;

    // Slashings.
    process_slashings(state, validator_statuses.total_balances.current_epoch, spec)?;

    // Final updates.
    process_final_updates(state, spec)?;

    // Rotate the epoch caches to suit the epoch transition.
    state.advance_caches();

    Ok(())
}

/// Update the following fields on the `BeaconState`:
///
/// - `justification_bitfield`.
/// - `previous_justified_epoch`
/// - `previous_justified_root`
/// - `current_justified_epoch`
/// - `current_justified_root`
/// - `finalized_epoch`
/// - `finalized_root`
///
/// Spec v0.6.1
pub fn process_justification_and_finalization<T: EthSpec>(
    state: &mut BeaconState<T>,
    total_balances: &TotalBalances,
    spec: &ChainSpec,
) -> Result<(), Error> {
    if state.current_epoch(spec) == spec.genesis_epoch {
        return Ok(());
    }

    let previous_epoch = state.previous_epoch(spec);
    let current_epoch = state.current_epoch(spec);

    let old_previous_justified_epoch = state.previous_justified_epoch;
    let old_current_justified_epoch = state.current_justified_epoch;

    // Process justifications
    state.previous_justified_epoch = state.current_justified_epoch;
    state.previous_justified_root = state.current_justified_root;
    state.justification_bitfield <<= 1;

    if total_balances.previous_epoch_target_attesters * 3 >= total_balances.previous_epoch * 2 {
        state.current_justified_epoch = previous_epoch;
        state.current_justified_root =
            *state.get_block_root_at_epoch(state.current_justified_epoch, spec)?;
        state.justification_bitfield |= 2;
    }
    // If the current epoch gets justified, fill the last bit.
    if total_balances.current_epoch_target_attesters * 3 >= total_balances.current_epoch * 2 {
        state.current_justified_epoch = current_epoch;
        state.current_justified_root =
            *state.get_block_root_at_epoch(state.current_justified_epoch, spec)?;
        state.justification_bitfield |= 1;
    }

    let bitfield = state.justification_bitfield;

    // The 2nd/3rd/4th most recent epochs are all justified, the 2nd using the 4th as source.
    if (bitfield >> 1) % 8 == 0b111 && old_previous_justified_epoch == current_epoch - 3 {
        state.finalized_epoch = old_previous_justified_epoch;
        state.finalized_root = *state.get_block_root_at_epoch(state.finalized_epoch, spec)?;
    }
    // The 2nd/3rd most recent epochs are both justified, the 2nd using the 3rd as source.
    if (bitfield >> 1) % 4 == 0b11 && state.previous_justified_epoch == current_epoch - 2 {
        state.finalized_epoch = old_previous_justified_epoch;
        state.finalized_root = *state.get_block_root_at_epoch(state.finalized_epoch, spec)?;
    }
    // The 1st/2nd/3rd most recent epochs are all justified, the 1st using the 2nd as source.
    if bitfield % 8 == 0b111 && state.current_justified_epoch == current_epoch - 2 {
        state.finalized_epoch = old_current_justified_epoch;
        state.finalized_root = *state.get_block_root_at_epoch(state.finalized_epoch, spec)?;
    }
    // The 1st/2nd most recent epochs are both justified, the 1st using the 2nd as source.
    if bitfield % 4 == 0b11 && state.current_justified_epoch == current_epoch - 1 {
        state.finalized_epoch = old_current_justified_epoch;
        state.finalized_root = *state.get_block_root_at_epoch(state.finalized_epoch, spec)?;
    }

    Ok(())
}

/// Updates the following fields on the `BeaconState`:
///
/// - `previous_crosslinks`
/// - `current_crosslinks`
///
/// Also returns a `WinningRootHashSet` for later use during epoch processing.
///
/// Spec v0.6.1
pub fn process_crosslinks<T: EthSpec>(
    state: &mut BeaconState<T>,
    spec: &ChainSpec,
) -> Result<WinningRootHashSet, Error> {
    let mut winning_root_for_shards: WinningRootHashSet = HashMap::new();

    state.previous_crosslinks = state.current_crosslinks.clone();

    for epoch in vec![state.previous_epoch(spec), state.current_epoch(spec)] {
        for offset in 0..state.get_epoch_committee_count(epoch, spec) {
            let shard = (state.get_epoch_start_shard(epoch, spec) + offset) % spec.shard_count;
            let crosslink_committee = state.get_crosslink_committee(epoch, shard, spec)?;

            let winning_root = winning_root(state, shard, epoch, spec)?;

            if let Some(winning_root) = winning_root {
                let total_committee_balance =
                    state.get_total_balance(&crosslink_committee.committee, spec)?;

                if 3 * winning_root.total_attesting_balance >= 2 * total_committee_balance {
                    state.current_crosslinks[shard as usize] = winning_root.crosslink.clone();
                }
                winning_root_for_shards.insert(shard, winning_root);
            }
        }
    }

    Ok(winning_root_for_shards)
}

/// Finish up an epoch update.
///
/// Spec v0.6.1
pub fn process_final_updates<T: EthSpec>(
    state: &mut BeaconState<T>,
    spec: &ChainSpec,
) -> Result<(), Error> {
    let current_epoch = state.current_epoch(spec);
    let next_epoch = state.next_epoch(spec);

    // Reset eth1 data votes.
    if (state.slot + 1) % spec.slots_per_eth1_voting_period == 0 {
        state.eth1_data_votes = vec![];
    }

    // Update effective balances with hysteresis (lag).
    for (index, validator) in state.validator_registry.iter_mut().enumerate() {
        let balance = state.balances[index];
        let half_increment = spec.effective_balance_increment / 2;
        if balance < validator.effective_balance
            || validator.effective_balance + 3 * half_increment < balance
        {
            validator.effective_balance = std::cmp::min(
                balance - balance % spec.effective_balance_increment,
                spec.max_effective_balance,
            );
        }
    }

    // Update start shard.
    state.latest_start_shard = (state.latest_start_shard
        + state.get_shard_delta(current_epoch, spec))
        % T::ShardCount::to_u64();

    // This is a hack to allow us to update index roots and slashed balances for the next epoch.
    //
    // The indentation here is to make it obvious where the weird stuff happens.
    {
        state.slot += 1;

        // Set active index root
        let active_index_root = Hash256::from_slice(
            &state
                .get_active_validator_indices(next_epoch + spec.activation_exit_delay)
                .tree_hash_root()[..],
        );
        state.set_active_index_root(next_epoch, active_index_root, spec)?;

        // Set total slashed balances
        state.set_slashed_balance(next_epoch, state.get_slashed_balance(current_epoch)?)?;

        // Set randao mix
        state.set_randao_mix(
            next_epoch,
            *state.get_randao_mix(current_epoch, spec)?,
            spec,
        )?;

        state.slot -= 1;
    }

    if next_epoch.as_u64() % (T::SlotsPerHistoricalRoot::to_u64() / spec.slots_per_epoch) == 0 {
        let historical_batch = state.historical_batch();
        state
            .historical_roots
            .push(Hash256::from_slice(&historical_batch.tree_hash_root()[..]));
    }

    // Rotate current/previous epoch attestations
    state.previous_epoch_attestations =
        std::mem::replace(&mut state.current_epoch_attestations, vec![]);

    Ok(())
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate rocket;

mod database;
mod domain;

use std::env;
use std::path::Path;

use aws_sdk_s3::types::ByteStream;
use aws_sdk_s3::Client;
use chrono::{DateTime, Datelike, SecondsFormat, Utc};
use dotenv::dotenv;
use rocket::data::ToByteUnit;
use rocket::http::ContentType;
use rocket::outcome::{try_outcome, Outcome};
use rocket::request::FromRequest;
use rocket::request::{self, FromParam};
use rocket::serde::json::Json;
use rocket::serde::Serialize;
use rocket::{Data, Request};
use rocket_db_pools::{Connection, Database};
use tokio::fs;

use database::Db;
use domain::entities::{ApiKeyId, DropId, User, UserId};
use domain::ApiKeyValue;

impl<'a> FromParam<'a> for DropId {
    type Error = String;

    fn from_param(param: &'a str) -> Result<Self, Self::Error> {
        DropId::try_from(param)
            .map_err(|err| format!("failed to parse '{}' as a drop ID: {}", param, err))
    }
}

impl<'a> FromParam<'a> for UserId {
    type Error = String;

    fn from_param(param: &'a str) -> Result<Self, Self::Error> {
        UserId::try_from(param)
            .map_err(|err| format!("failed to parse '{}' as a user ID: {}", param, err))
    }
}

pub struct ApiKeyBearer {
    user: User,
}

struct ApiKeyLookup {
    pub id: String,
    pub username: String,
    pub full_name: Option<String>,
}

#[rocket::async_trait]
impl<'r> FromRequest<'r> for ApiKeyBearer {
    type Error = String;

    async fn from_request(req: &'r Request<'_>) -> request::Outcome<Self, Self::Error> {
        let mut db = try_outcome!(req
            .guard::<Connection<Db>>()
            .await
            .map_failure(|f| (f.0, "failed to acquire database connection".to_string())));

        match req.headers().get_one("Authorization") {
            Some(authorization) => {
                let token = authorization.replace("Bearer ", "");

                let api_key = sqlx::query_as!(
                    ApiKeyLookup,
                    "
                    select user.id, user.username, user.full_name from api_key
                    inner join user on user.id = api_key.user_id
                    where api_key.value = $1
                    ",
                    token
                )
                .fetch_one(&mut *db)
                .await
                .unwrap();

                Outcome::Success(ApiKeyBearer {
                    user: User {
                        id: UserId::try_from(api_key.id).unwrap(),
                        username: api_key.username,
                        full_name: api_key.full_name,
                    },
                })
            }
            _ => Outcome::Forward(()),
        }
    }
}

#[get("/")]
async fn index() -> &'static str {
    "滴"
}

#[post("/drops", data = "<drop>")]
async fn upload_drop(_bearer: ApiKeyBearer, drop: Data<'_>) -> std::io::Result<String> {
    let id = DropId::new();

    let temp_dir = Path::new("temp");
    fs::create_dir_all(temp_dir).await.unwrap();

    let temp_filepath = temp_dir.join(id.to_string());

    drop.open(25.mebibytes()).into_file(&temp_filepath).await?;

    let s3_config = aws_config::load_from_env().await;
    let client = Client::new(&s3_config);

    let body = ByteStream::from_path(&temp_filepath).await?;

    client
        .put_object()
        .bucket(env::var("AWS_S3_BUCKET").unwrap())
        .key(format!(
            "drops/{}/{}",
            id.created_at().date().year(),
            id.unprefixed()
        ))
        .body(body)
        .send()
        .await
        .expect("failed to upload to S3");

    tokio::fs::remove_file(temp_filepath).await.ok();

    Ok(id.to_string())
}

#[get("/drops/<id>")]
async fn get_drop(id: DropId) -> (ContentType, Vec<u8>) {
    let s3_config = aws_config::load_from_env().await;
    let client = Client::new(&s3_config);

    let object = client
        .get_object()
        .bucket(env::var("AWS_S3_BUCKET").unwrap())
        .key(format!(
            "drops/{}/{}",
            id.created_at().date().year(),
            id.unprefixed()
        ))
        .send()
        .await
        .expect("failed to retrieve object from S3");

    let bytes = object
        .body
        .collect()
        .await
        .map(|data| data.into_bytes())
        .unwrap()
        .to_vec();

    let mime_type = tree_magic_mini::from_u8(&bytes);

    let content_type = ContentType::parse_flexible(mime_type).unwrap_or(ContentType::Binary);

    (content_type, bytes)
}

#[derive(Debug, Serialize)]
pub struct DropMetadata {
    pub created_at: DateTime<Utc>,
    pub mime_type: String,
    pub size_in_bytes: i64,
}

#[get("/drops/<id>/meta")]
async fn get_drop_metadata(id: DropId) -> Json<DropMetadata> {
    let s3_config = aws_config::load_from_env().await;
    let client = Client::new(&s3_config);

    let object = client
        .get_object()
        .bucket(env::var("AWS_S3_BUCKET").unwrap())
        .key(format!(
            "drops/{}/{}",
            id.created_at().date().year(),
            id.unprefixed()
        ))
        .send()
        .await
        .expect("failed to retrieve object from S3");

    let bytes = object
        .body
        .collect()
        .await
        .map(|data| data.into_bytes())
        .unwrap()
        .to_vec();

    let mime_type = tree_magic_mini::from_u8(&bytes).to_string();

    Json(DropMetadata {
        created_at: id.created_at(),
        mime_type,
        size_in_bytes: object.content_length,
    })
}

#[derive(Debug, Serialize)]
pub struct UserJson {
    pub id: String,
    pub username: String,
    pub full_name: Option<String>,
}

#[post("/users")]
async fn create_user(
    _bearer: ApiKeyBearer,
    mut db: Connection<Db>,
) -> database::Result<Json<UserJson>> {
    let id = UserId::new();

    let username = "maxdeviant";
    let full_name = "Marshall Bowers";

    let unprefixed_id = id.unprefixed();
    let now = Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true);

    sqlx::query!(
        "
        insert into user (id, created_at, updated_at, username, full_name)
        values (?, ?, ?, ?, ?)
        ",
        unprefixed_id,
        now,
        now,
        username,
        full_name
    )
    .execute(&mut *db)
    .await?;

    Ok(Json(UserJson {
        id: id.to_string(),
        username: username.to_string(),
        full_name: Some(full_name.to_string()),
    }))
}

#[derive(Debug, Serialize)]
pub struct ApiKeyJson {
    pub id: String,
    pub value: String,
}

#[post("/keys")]
async fn generate_api_key(
    bearer: ApiKeyBearer,
    mut db: Connection<Db>,
) -> database::Result<Json<ApiKeyJson>> {
    let id = ApiKeyId::new();
    let value = ApiKeyValue::new();

    let unprefixed_id = id.unprefixed();
    let unprefixed_user_id = bearer.user.id.unprefixed();
    let key_value = value.to_string();
    let now = Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true);

    sqlx::query!(
        "
        insert into api_key (id, created_at, updated_at, user_id, value)
        values (?, ?, ?, ?, ?)
        ",
        unprefixed_id,
        now,
        now,
        unprefixed_user_id,
        key_value
    )
    .execute(&mut *db)
    .await?;

    Ok(Json(ApiKeyJson {
        id: id.to_string(),
        value: value.to_string(),
    }))
}

async fn ensure_root_user_exists() {
    use sqlx::{Connection, SqliteConnection};

    let root_username = env::var("ROOT_USERNAME").unwrap();
    let root_api_key = env::var("ROOT_API_KEY").unwrap();

    let database_url = env::var("DATABASE_URL").unwrap();
    let mut connection = SqliteConnection::connect(&database_url).await.unwrap();

    let user_id = UserId::new();
    let unprefixed_user_id = user_id.unprefixed();
    let now = Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true);

    #[derive(Debug)]
    struct RootUser {
        id: String,
    }

    let root_user = sqlx::query_as!(
        RootUser,
        "
        insert into user (id, created_at, updated_at, username)
        values (?, ?, ?, ?)
        -- This is a no-op update just so we can get the `id` returned.
        on conflict (username) do update set updated_at = updated_at
        returning id
        ",
        unprefixed_user_id,
        now,
        now,
        root_username
    )
    .fetch_one(&mut connection)
    .await
    .expect("failed to insert root user");

    println!("Root user: {:?}", root_user);

    let key_id = ApiKeyId::new();
    let unprefixed_key_id = key_id.unprefixed();
    let now = Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true);

    sqlx::query!(
        "
        insert into api_key (id, created_at, updated_at, user_id, value)
        values (?, ?, ?, ?, ?)
        on conflict (value) do nothing
        ",
        unprefixed_key_id,
        now,
        now,
        root_user.id,
        root_api_key
    )
    .execute(&mut connection)
    .await
    .unwrap();
}

#[rocket::main]
async fn main() -> Result<(), rocket::Error> {
    dotenv().ok();

    ensure_root_user_exists().await;

    let rocket = rocket::build()
        .mount(
            "/",
            routes![index, upload_drop, get_drop, get_drop_metadata,],
        )
        .mount("/x/", routes![create_user, generate_api_key])
        .attach(Db::init())
        .ignite()
        .await?;

    let _ = rocket.launch().await?;

    Ok(())
}
<commit_msg>Set caching headers on drop responses<commit_after>#[macro_use]
extern crate rocket;

mod database;
mod domain;

use std::env;
use std::path::Path;

use aws_sdk_s3::types::ByteStream;
use aws_sdk_s3::Client;
use chrono::{DateTime, Datelike, SecondsFormat, Utc};
use dotenv::dotenv;
use rocket::data::ToByteUnit;
use rocket::http::hyper::header;
use rocket::http::{ContentType, Header};
use rocket::outcome::{try_outcome, Outcome};
use rocket::request::FromRequest;
use rocket::request::{self, FromParam};
use rocket::serde::json::Json;
use rocket::serde::Serialize;
use rocket::{Data, Request};
use rocket_db_pools::{Connection, Database};
use tokio::fs;

use database::Db;
use domain::entities::{ApiKeyId, DropId, User, UserId};
use domain::ApiKeyValue;

impl<'a> FromParam<'a> for DropId {
    type Error = String;

    fn from_param(param: &'a str) -> Result<Self, Self::Error> {
        DropId::try_from(param)
            .map_err(|err| format!("failed to parse '{}' as a drop ID: {}", param, err))
    }
}

impl<'a> FromParam<'a> for UserId {
    type Error = String;

    fn from_param(param: &'a str) -> Result<Self, Self::Error> {
        UserId::try_from(param)
            .map_err(|err| format!("failed to parse '{}' as a user ID: {}", param, err))
    }
}

pub struct ApiKeyBearer {
    user: User,
}

struct ApiKeyLookup {
    pub id: String,
    pub username: String,
    pub full_name: Option<String>,
}

#[rocket::async_trait]
impl<'r> FromRequest<'r> for ApiKeyBearer {
    type Error = String;

    async fn from_request(req: &'r Request<'_>) -> request::Outcome<Self, Self::Error> {
        let mut db = try_outcome!(req
            .guard::<Connection<Db>>()
            .await
            .map_failure(|f| (f.0, "failed to acquire database connection".to_string())));

        match req.headers().get_one("Authorization") {
            Some(authorization) => {
                let token = authorization.replace("Bearer ", "");

                let api_key = sqlx::query_as!(
                    ApiKeyLookup,
                    "
                    select user.id, user.username, user.full_name from api_key
                    inner join user on user.id = api_key.user_id
                    where api_key.value = $1
                    ",
                    token
                )
                .fetch_one(&mut *db)
                .await
                .unwrap();

                Outcome::Success(ApiKeyBearer {
                    user: User {
                        id: UserId::try_from(api_key.id).unwrap(),
                        username: api_key.username,
                        full_name: api_key.full_name,
                    },
                })
            }
            _ => Outcome::Forward(()),
        }
    }
}

#[get("/")]
async fn index() -> &'static str {
    "滴"
}

#[post("/drops", data = "<drop>")]
async fn upload_drop(_bearer: ApiKeyBearer, drop: Data<'_>) -> std::io::Result<String> {
    let id = DropId::new();

    let temp_dir = Path::new("temp");
    fs::create_dir_all(temp_dir).await.unwrap();

    let temp_filepath = temp_dir.join(id.to_string());

    drop.open(25.mebibytes()).into_file(&temp_filepath).await?;

    let s3_config = aws_config::load_from_env().await;
    let client = Client::new(&s3_config);

    let body = ByteStream::from_path(&temp_filepath).await?;

    client
        .put_object()
        .bucket(env::var("AWS_S3_BUCKET").unwrap())
        .key(format!(
            "drops/{}/{}",
            id.created_at().date().year(),
            id.unprefixed()
        ))
        .body(body)
        .send()
        .await
        .expect("failed to upload to S3");

    tokio::fs::remove_file(temp_filepath).await.ok();

    Ok(id.to_string())
}

#[derive(Responder)]
struct DropResponse {
    bytes: Vec<u8>,
    content_type: ContentType,
    cache_control: Header<'static>,
    e_tag: Header<'static>,
}

#[get("/drops/<id>")]
async fn get_drop(id: DropId) -> DropResponse {
    let s3_config = aws_config::load_from_env().await;
    let client = Client::new(&s3_config);

    let object = client
        .get_object()
        .bucket(env::var("AWS_S3_BUCKET").unwrap())
        .key(format!(
            "drops/{}/{}",
            id.created_at().date().year(),
            id.unprefixed()
        ))
        .send()
        .await
        .expect("failed to retrieve object from S3");

    let e_tag = object.e_tag().to_owned().unwrap_or("").to_owned();

    let bytes = object
        .body
        .collect()
        .await
        .map(|data| data.into_bytes())
        .unwrap()
        .to_vec();

    let mime_type = tree_magic_mini::from_u8(&bytes);

    let content_type = ContentType::parse_flexible(mime_type).unwrap_or(ContentType::Binary);

    DropResponse {
        bytes,
        content_type,
        cache_control: Header::new(header::CACHE_CONTROL.to_string(), "public, max-age=3600"),
        e_tag: Header::new(header::ETAG.to_string(), e_tag),
    }
}

#[derive(Debug, Serialize)]
pub struct DropMetadata {
    pub created_at: DateTime<Utc>,
    pub mime_type: String,
    pub size_in_bytes: i64,
}

#[get("/drops/<id>/meta")]
async fn get_drop_metadata(id: DropId) -> Json<DropMetadata> {
    let s3_config = aws_config::load_from_env().await;
    let client = Client::new(&s3_config);

    let object = client
        .get_object()
        .bucket(env::var("AWS_S3_BUCKET").unwrap())
        .key(format!(
            "drops/{}/{}",
            id.created_at().date().year(),
            id.unprefixed()
        ))
        .send()
        .await
        .expect("failed to retrieve object from S3");

    let bytes = object
        .body
        .collect()
        .await
        .map(|data| data.into_bytes())
        .unwrap()
        .to_vec();

    let mime_type = tree_magic_mini::from_u8(&bytes).to_string();

    Json(DropMetadata {
        created_at: id.created_at(),
        mime_type,
        size_in_bytes: object.content_length,
    })
}

#[derive(Debug, Serialize)]
pub struct UserJson {
    pub id: String,
    pub username: String,
    pub full_name: Option<String>,
}

#[post("/users")]
async fn create_user(
    _bearer: ApiKeyBearer,
    mut db: Connection<Db>,
) -> database::Result<Json<UserJson>> {
    let id = UserId::new();

    let username = "maxdeviant";
    let full_name = "Marshall Bowers";

    let unprefixed_id = id.unprefixed();
    let now = Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true);

    sqlx::query!(
        "
        insert into user (id, created_at, updated_at, username, full_name)
        values (?, ?, ?, ?, ?)
        ",
        unprefixed_id,
        now,
        now,
        username,
        full_name
    )
    .execute(&mut *db)
    .await?;

    Ok(Json(UserJson {
        id: id.to_string(),
        username: username.to_string(),
        full_name: Some(full_name.to_string()),
    }))
}

#[derive(Debug, Serialize)]
pub struct ApiKeyJson {
    pub id: String,
    pub value: String,
}

#[post("/keys")]
async fn generate_api_key(
    bearer: ApiKeyBearer,
    mut db: Connection<Db>,
) -> database::Result<Json<ApiKeyJson>> {
    let id = ApiKeyId::new();
    let value = ApiKeyValue::new();

    let unprefixed_id = id.unprefixed();
    let unprefixed_user_id = bearer.user.id.unprefixed();
    let key_value = value.to_string();
    let now = Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true);

    sqlx::query!(
        "
        insert into api_key (id, created_at, updated_at, user_id, value)
        values (?, ?, ?, ?, ?)
        ",
        unprefixed_id,
        now,
        now,
        unprefixed_user_id,
        key_value
    )
    .execute(&mut *db)
    .await?;

    Ok(Json(ApiKeyJson {
        id: id.to_string(),
        value: value.to_string(),
    }))
}

async fn ensure_root_user_exists() {
    use sqlx::{Connection, SqliteConnection};

    let root_username = env::var("ROOT_USERNAME").unwrap();
    let root_api_key = env::var("ROOT_API_KEY").unwrap();

    let database_url = env::var("DATABASE_URL").unwrap();
    let mut connection = SqliteConnection::connect(&database_url).await.unwrap();

    let user_id = UserId::new();
    let unprefixed_user_id = user_id.unprefixed();
    let now = Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true);

    #[derive(Debug)]
    struct RootUser {
        id: String,
    }

    let root_user = sqlx::query_as!(
        RootUser,
        "
        insert into user (id, created_at, updated_at, username)
        values (?, ?, ?, ?)
        -- This is a no-op update just so we can get the `id` returned.
        on conflict (username) do update set updated_at = updated_at
        returning id
        ",
        unprefixed_user_id,
        now,
        now,
        root_username
    )
    .fetch_one(&mut connection)
    .await
    .expect("failed to insert root user");

    println!("Root user: {:?}", root_user);

    let key_id = ApiKeyId::new();
    let unprefixed_key_id = key_id.unprefixed();
    let now = Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true);

    sqlx::query!(
        "
        insert into api_key (id, created_at, updated_at, user_id, value)
        values (?, ?, ?, ?, ?)
        on conflict (value) do nothing
        ",
        unprefixed_key_id,
        now,
        now,
        root_user.id,
        root_api_key
    )
    .execute(&mut connection)
    .await
    .unwrap();
}

#[rocket::main]
async fn main() -> Result<(), rocket::Error> {
    dotenv().ok();

    ensure_root_user_exists().await;

    let rocket = rocket::build()
        .mount(
            "/",
            routes![index, upload_drop, get_drop, get_drop_metadata,],
        )
        .mount("/x/", routes![create_user, generate_api_key])
        .attach(Db::init())
        .ignite()
        .await?;

    let _ = rocket.launch().await?;

    Ok(())
}
<|endoftext|>


<commit_before>#[derive(Eq, PartialEq)]
enum Direction {
    North,
    East,
    South,
    West,
    Empty,
}

struct Table {
    w: i32,
    h: i32,
}

struct Robot {
    x: i32,
    y: i32,
    d: Direction,
    t: Table,
}

impl Direction {
    fn string(&self) -> String {
        let s = match *self {
            Direction::North => "North",
            Direction::East => "East",
            Direction::South => "South",
            Direction::West => "West",
            _ => "Empty",
        };

        s.to_string()
    }
}

impl Default for Robot {
    fn default() -> Robot {
        let table = Table{w: 5, h: 5};

        Robot{x: 0, y: 0, d: Direction::Empty, t: table}
    }
}

impl Table {
    fn valid_pos(&self, x: i32, y: i32) -> bool {
        x >= 0 && x < self.w && y >= 0 && y < self.h
    }
}

impl Robot {
    fn place(&mut self, x: i32, y: i32, d: Direction) {
        if self.t.valid_pos(x, y) {
            self.x = x;
            self.y = y;
            self.d = d;
        }
    }

    fn right(&mut self) {
        let d = match self.d {
            Direction::North => Direction::East,
            Direction::East => Direction::South,
            Direction::South => Direction::West,
            Direction::West => Direction::North,
            _ => Direction::Empty,
        };

        self.d = d;
    }

    fn left(&mut self) {
        let d = match self.d {
            Direction::North => Direction::West,
            Direction::West => Direction::South,
            Direction::South => Direction::East,
            Direction::East => Direction::North,
            _ => Direction::Empty,
        };

        self.d = d;
    }

    fn step(&mut self) {
        match self.d {
            Direction::North => if self.t.valid_pos(self.x, self.y - 1) {self.y -= 1;},
            Direction::East => if self.t.valid_pos(self.x + 1, self.y) {self.x += 1;},
            Direction::South => if self.t.valid_pos(self.x, self.y + 1) {self.y += 1;},
            Direction::West => if self.t.valid_pos(self.x - 1, self.y) {self.y -= 1;},
            _ => (),
        }
    }

    fn report(&self) {
        println!("{}, {}, {}", self.x, self.y, self.d.string());
    }
}

fn main() {
    let mut robot: Robot = Default::default();

    robot.place(1, 2, Direction::North);
    robot.left();
}
<commit_msg>Engine working. Just need parser<commit_after>#[derive(Eq, PartialEq)]
enum Direction {
    North,
    East,
    South,
    West,
    Empty,
}

struct Table {
    w: i32,
    h: i32,
}

struct Robot {
    x: i32,
    y: i32,
    d: Direction,
    t: Table,
}

impl Direction {
    fn string(&self) -> String {
        let s = match *self {
            Direction::North => "North",
            Direction::East => "East",
            Direction::South => "South",
            Direction::West => "West",
            _ => "Empty",
        };

        s.to_string()
    }
}

impl Default for Robot {
    fn default() -> Robot {
        let table = Table{w: 5, h: 5};

        Robot{x: 0, y: 0, d: Direction::Empty, t: table}
    }
}

impl Table {
    fn valid_pos(&self, x: i32, y: i32) -> bool {
        x >= 0 && x < self.w && y >= 0 && y < self.h
    }
}

impl Robot {
    fn place(&mut self, x: i32, y: i32, d: Direction) {
        if self.t.valid_pos(x, y) {
            self.x = x;
            self.y = y;
            self.d = d;
        }
    }

    fn right(&mut self) {
        let d = match self.d {
            Direction::North => Direction::East,
            Direction::East => Direction::South,
            Direction::South => Direction::West,
            Direction::West => Direction::North,
            _ => Direction::Empty,
        };

        self.d = d;
    }

    fn left(&mut self) {
        let d = match self.d {
            Direction::North => Direction::West,
            Direction::West => Direction::South,
            Direction::South => Direction::East,
            Direction::East => Direction::North,
            _ => Direction::Empty,
        };

        self.d = d;
    }

    fn step(&mut self) {
        match self.d {
            Direction::North => if self.t.valid_pos(self.x, self.y + 1) {self.y += 1;},
            Direction::East => if self.t.valid_pos(self.x + 1, self.y) {self.x += 1;},
            Direction::South => if self.t.valid_pos(self.x, self.y - 1) {self.y -= 1;},
            Direction::West => if self.t.valid_pos(self.x - 1, self.y) {self.y -= 1;},
            _ => (),
        }
    }

    fn report(&self) {
        println!("{}, {}, {}", self.x, self.y, self.d.string());
    }
}

fn main() {
    let mut robot: Robot = Default::default();
}
<|endoftext|>


<commit_before>use core::sync::atomic::Ordering::SeqCst;
use core::sync::atomic::*;

#[test]
fn bool_() {
    let a = AtomicBool::new(false);
    assert_eq!(a.compare_exchange(false, true, SeqCst, SeqCst), Ok(false));
    assert_eq!(a.compare_exchange(false, true, SeqCst, SeqCst), Err(true));

    a.store(false, SeqCst);
    assert_eq!(a.compare_exchange(false, true, SeqCst, SeqCst), Ok(false));
}

#[test]
fn bool_and() {
    let a = AtomicBool::new(true);
    assert_eq!(a.fetch_and(false, SeqCst), true);
    assert_eq!(a.load(SeqCst), false);
}

#[test]
fn bool_nand() {
    let a = AtomicBool::new(false);
    assert_eq!(a.fetch_nand(false, SeqCst), false);
    assert_eq!(a.load(SeqCst), true);
    assert_eq!(a.fetch_nand(false, SeqCst), true);
    assert_eq!(a.load(SeqCst), true);
    assert_eq!(a.fetch_nand(true, SeqCst), true);
    assert_eq!(a.load(SeqCst), false);
    assert_eq!(a.fetch_nand(true, SeqCst), false);
    assert_eq!(a.load(SeqCst), true);
}

#[test]
fn uint_and() {
    let x = AtomicUsize::new(0xf731);
    assert_eq!(x.fetch_and(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 & 0x137f);
}

#[test]
fn uint_nand() {
    let x = AtomicUsize::new(0xf731);
    assert_eq!(x.fetch_nand(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), !(0xf731 & 0x137f));
}

#[test]
fn uint_or() {
    let x = AtomicUsize::new(0xf731);
    assert_eq!(x.fetch_or(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 | 0x137f);
}

#[test]
fn uint_xor() {
    let x = AtomicUsize::new(0xf731);
    assert_eq!(x.fetch_xor(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 ^ 0x137f);
}

#[test]
fn int_and() {
    let x = AtomicIsize::new(0xf731);
    assert_eq!(x.fetch_and(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 & 0x137f);
}

#[test]
fn int_nand() {
    let x = AtomicIsize::new(0xf731);
    assert_eq!(x.fetch_nand(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), !(0xf731 & 0x137f));
}

#[test]
fn int_or() {
    let x = AtomicIsize::new(0xf731);
    assert_eq!(x.fetch_or(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 | 0x137f);
}

#[test]
fn int_xor() {
    let x = AtomicIsize::new(0xf731);
    assert_eq!(x.fetch_xor(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 ^ 0x137f);
}

static S_FALSE: AtomicBool = AtomicBool::new(false);
static S_TRUE: AtomicBool = AtomicBool::new(true);
static S_INT: AtomicIsize = AtomicIsize::new(0);
static S_UINT: AtomicUsize = AtomicUsize::new(0);

#[test]
fn static_init() {
    // Note that we're not really testing the mutability here but it's important
    // on Android at the moment (#49775)
    assert!(!S_FALSE.swap(true, SeqCst));
    assert!(S_TRUE.swap(false, SeqCst));
    assert!(S_INT.fetch_add(1, SeqCst) == 0);
    assert!(S_UINT.fetch_add(1, SeqCst) == 0);
}

#[test]
fn atomic_access_bool() {
    static mut ATOMIC: AtomicBool = AtomicBool::new(false);

    unsafe {
        assert_eq!(*ATOMIC.get_mut(), false);
        ATOMIC.store(true, SeqCst);
        assert_eq!(*ATOMIC.get_mut(), true);
        ATOMIC.fetch_or(false, SeqCst);
        assert_eq!(*ATOMIC.get_mut(), true);
        ATOMIC.fetch_and(false, SeqCst);
        assert_eq!(*ATOMIC.get_mut(), false);
        ATOMIC.fetch_nand(true, SeqCst);
        assert_eq!(*ATOMIC.get_mut(), true);
        ATOMIC.fetch_xor(true, SeqCst);
        assert_eq!(*ATOMIC.get_mut(), false);
    }
}

#[test]
fn atomic_alignment() {
    use std::mem::{align_of, size_of};

    #[cfg(target_has_atomic = "8")]
    assert_eq!(align_of::<AtomicBool>(), size_of::<AtomicBool>());
    #[cfg(target_has_atomic = "ptr")]
    assert_eq!(align_of::<AtomicPtr<u8>>(), size_of::<AtomicPtr<u8>>());
    #[cfg(target_has_atomic = "8")]
    assert_eq!(align_of::<AtomicU8>(), size_of::<AtomicU8>());
    #[cfg(target_has_atomic = "8")]
    assert_eq!(align_of::<AtomicI8>(), size_of::<AtomicI8>());
    #[cfg(target_has_atomic = "16")]
    assert_eq!(align_of::<AtomicU16>(), size_of::<AtomicU16>());
    #[cfg(target_has_atomic = "16")]
    assert_eq!(align_of::<AtomicI16>(), size_of::<AtomicI16>());
    #[cfg(target_has_atomic = "32")]
    assert_eq!(align_of::<AtomicU32>(), size_of::<AtomicU32>());
    #[cfg(target_has_atomic = "32")]
    assert_eq!(align_of::<AtomicI32>(), size_of::<AtomicI32>());
    #[cfg(target_has_atomic = "64")]
    assert_eq!(align_of::<AtomicU64>(), size_of::<AtomicU64>());
    #[cfg(target_has_atomic = "64")]
    assert_eq!(align_of::<AtomicI64>(), size_of::<AtomicI64>());
    #[cfg(target_has_atomic = "128")]
    assert_eq!(align_of::<AtomicU128>(), size_of::<AtomicU128>());
    #[cfg(target_has_atomic = "128")]
    assert_eq!(align_of::<AtomicI128>(), size_of::<AtomicI128>());
    #[cfg(target_has_atomic = "ptr")]
    assert_eq!(align_of::<AtomicUsize>(), size_of::<AtomicUsize>());
    #[cfg(target_has_atomic = "ptr")]
    assert_eq!(align_of::<AtomicIsize>(), size_of::<AtomicIsize>());
}

#[test]
fn atomic_compare_exchange() {
    use Ordering::*;

    static ATOMIC: AtomicIsize = AtomicIsize::new(0);

    ATOMIC.compare_exchange(0, 1, Relaxed, Relaxed).ok();
    ATOMIC.compare_exchange(0, 1, Acquire, Relaxed).ok();
    ATOMIC.compare_exchange(0, 1, Release, Relaxed).ok();
    ATOMIC.compare_exchange(0, 1, AcqRel, Relaxed).ok();
    ATOMIC.compare_exchange(0, 1, SeqCst, Relaxed).ok();
    ATOMIC.compare_exchange(0, 1, Acquire, Acquire).ok();
    ATOMIC.compare_exchange(0, 1, AcqRel, Acquire).ok();
    ATOMIC.compare_exchange(0, 1, SeqCst, Acquire).ok();
    ATOMIC.compare_exchange(0, 1, SeqCst, SeqCst).ok();
    ATOMIC.compare_exchange_weak(0, 1, Relaxed, Relaxed).ok();
    ATOMIC.compare_exchange_weak(0, 1, Acquire, Relaxed).ok();
    ATOMIC.compare_exchange_weak(0, 1, Release, Relaxed).ok();
    ATOMIC.compare_exchange_weak(0, 1, AcqRel, Relaxed).ok();
    ATOMIC.compare_exchange_weak(0, 1, SeqCst, Relaxed).ok();
    ATOMIC.compare_exchange_weak(0, 1, Acquire, Acquire).ok();
    ATOMIC.compare_exchange_weak(0, 1, AcqRel, Acquire).ok();
    ATOMIC.compare_exchange_weak(0, 1, SeqCst, Acquire).ok();
    ATOMIC.compare_exchange_weak(0, 1, SeqCst, SeqCst).ok();
}
<commit_msg>Add tests for Atomic*::fetch_{min,max}<commit_after>use core::sync::atomic::Ordering::SeqCst;
use core::sync::atomic::*;

#[test]
fn bool_() {
    let a = AtomicBool::new(false);
    assert_eq!(a.compare_exchange(false, true, SeqCst, SeqCst), Ok(false));
    assert_eq!(a.compare_exchange(false, true, SeqCst, SeqCst), Err(true));

    a.store(false, SeqCst);
    assert_eq!(a.compare_exchange(false, true, SeqCst, SeqCst), Ok(false));
}

#[test]
fn bool_and() {
    let a = AtomicBool::new(true);
    assert_eq!(a.fetch_and(false, SeqCst), true);
    assert_eq!(a.load(SeqCst), false);
}

#[test]
fn bool_nand() {
    let a = AtomicBool::new(false);
    assert_eq!(a.fetch_nand(false, SeqCst), false);
    assert_eq!(a.load(SeqCst), true);
    assert_eq!(a.fetch_nand(false, SeqCst), true);
    assert_eq!(a.load(SeqCst), true);
    assert_eq!(a.fetch_nand(true, SeqCst), true);
    assert_eq!(a.load(SeqCst), false);
    assert_eq!(a.fetch_nand(true, SeqCst), false);
    assert_eq!(a.load(SeqCst), true);
}

#[test]
fn uint_and() {
    let x = AtomicUsize::new(0xf731);
    assert_eq!(x.fetch_and(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 & 0x137f);
}

#[test]
fn uint_nand() {
    let x = AtomicUsize::new(0xf731);
    assert_eq!(x.fetch_nand(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), !(0xf731 & 0x137f));
}

#[test]
fn uint_or() {
    let x = AtomicUsize::new(0xf731);
    assert_eq!(x.fetch_or(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 | 0x137f);
}

#[test]
fn uint_xor() {
    let x = AtomicUsize::new(0xf731);
    assert_eq!(x.fetch_xor(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 ^ 0x137f);
}

#[test]
fn uint_min() {
    let x = AtomicUsize::new(0xf731);
    assert_eq!(x.fetch_min(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0x137f);
    assert_eq!(x.fetch_min(0xf731, SeqCst), 0x137f);
    assert_eq!(x.load(SeqCst), 0x137f);
}

#[test]
fn uint_max() {
    let x = AtomicUsize::new(0x137f);
    assert_eq!(x.fetch_max(0xf731, SeqCst), 0x137f);
    assert_eq!(x.load(SeqCst), 0xf731);
    assert_eq!(x.fetch_max(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731);
}

#[test]
fn int_and() {
    let x = AtomicIsize::new(0xf731);
    assert_eq!(x.fetch_and(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 & 0x137f);
}

#[test]
fn int_nand() {
    let x = AtomicIsize::new(0xf731);
    assert_eq!(x.fetch_nand(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), !(0xf731 & 0x137f));
}

#[test]
fn int_or() {
    let x = AtomicIsize::new(0xf731);
    assert_eq!(x.fetch_or(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 | 0x137f);
}

#[test]
fn int_xor() {
    let x = AtomicIsize::new(0xf731);
    assert_eq!(x.fetch_xor(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731 ^ 0x137f);
}

#[test]
fn int_min() {
    let x = AtomicIsize::new(0xf731);
    assert_eq!(x.fetch_min(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0x137f);
    assert_eq!(x.fetch_min(0xf731, SeqCst), 0x137f);
    assert_eq!(x.load(SeqCst), 0x137f);
}

#[test]
fn int_max() {
    let x = AtomicIsize::new(0x137f);
    assert_eq!(x.fetch_max(0xf731, SeqCst), 0x137f);
    assert_eq!(x.load(SeqCst), 0xf731);
    assert_eq!(x.fetch_max(0x137f, SeqCst), 0xf731);
    assert_eq!(x.load(SeqCst), 0xf731);
}

static S_FALSE: AtomicBool = AtomicBool::new(false);
static S_TRUE: AtomicBool = AtomicBool::new(true);
static S_INT: AtomicIsize = AtomicIsize::new(0);
static S_UINT: AtomicUsize = AtomicUsize::new(0);

#[test]
fn static_init() {
    // Note that we're not really testing the mutability here but it's important
    // on Android at the moment (#49775)
    assert!(!S_FALSE.swap(true, SeqCst));
    assert!(S_TRUE.swap(false, SeqCst));
    assert!(S_INT.fetch_add(1, SeqCst) == 0);
    assert!(S_UINT.fetch_add(1, SeqCst) == 0);
}

#[test]
fn atomic_access_bool() {
    static mut ATOMIC: AtomicBool = AtomicBool::new(false);

    unsafe {
        assert_eq!(*ATOMIC.get_mut(), false);
        ATOMIC.store(true, SeqCst);
        assert_eq!(*ATOMIC.get_mut(), true);
        ATOMIC.fetch_or(false, SeqCst);
        assert_eq!(*ATOMIC.get_mut(), true);
        ATOMIC.fetch_and(false, SeqCst);
        assert_eq!(*ATOMIC.get_mut(), false);
        ATOMIC.fetch_nand(true, SeqCst);
        assert_eq!(*ATOMIC.get_mut(), true);
        ATOMIC.fetch_xor(true, SeqCst);
        assert_eq!(*ATOMIC.get_mut(), false);
    }
}

#[test]
fn atomic_alignment() {
    use std::mem::{align_of, size_of};

    #[cfg(target_has_atomic = "8")]
    assert_eq!(align_of::<AtomicBool>(), size_of::<AtomicBool>());
    #[cfg(target_has_atomic = "ptr")]
    assert_eq!(align_of::<AtomicPtr<u8>>(), size_of::<AtomicPtr<u8>>());
    #[cfg(target_has_atomic = "8")]
    assert_eq!(align_of::<AtomicU8>(), size_of::<AtomicU8>());
    #[cfg(target_has_atomic = "8")]
    assert_eq!(align_of::<AtomicI8>(), size_of::<AtomicI8>());
    #[cfg(target_has_atomic = "16")]
    assert_eq!(align_of::<AtomicU16>(), size_of::<AtomicU16>());
    #[cfg(target_has_atomic = "16")]
    assert_eq!(align_of::<AtomicI16>(), size_of::<AtomicI16>());
    #[cfg(target_has_atomic = "32")]
    assert_eq!(align_of::<AtomicU32>(), size_of::<AtomicU32>());
    #[cfg(target_has_atomic = "32")]
    assert_eq!(align_of::<AtomicI32>(), size_of::<AtomicI32>());
    #[cfg(target_has_atomic = "64")]
    assert_eq!(align_of::<AtomicU64>(), size_of::<AtomicU64>());
    #[cfg(target_has_atomic = "64")]
    assert_eq!(align_of::<AtomicI64>(), size_of::<AtomicI64>());
    #[cfg(target_has_atomic = "128")]
    assert_eq!(align_of::<AtomicU128>(), size_of::<AtomicU128>());
    #[cfg(target_has_atomic = "128")]
    assert_eq!(align_of::<AtomicI128>(), size_of::<AtomicI128>());
    #[cfg(target_has_atomic = "ptr")]
    assert_eq!(align_of::<AtomicUsize>(), size_of::<AtomicUsize>());
    #[cfg(target_has_atomic = "ptr")]
    assert_eq!(align_of::<AtomicIsize>(), size_of::<AtomicIsize>());
}

#[test]
fn atomic_compare_exchange() {
    use Ordering::*;

    static ATOMIC: AtomicIsize = AtomicIsize::new(0);

    ATOMIC.compare_exchange(0, 1, Relaxed, Relaxed).ok();
    ATOMIC.compare_exchange(0, 1, Acquire, Relaxed).ok();
    ATOMIC.compare_exchange(0, 1, Release, Relaxed).ok();
    ATOMIC.compare_exchange(0, 1, AcqRel, Relaxed).ok();
    ATOMIC.compare_exchange(0, 1, SeqCst, Relaxed).ok();
    ATOMIC.compare_exchange(0, 1, Acquire, Acquire).ok();
    ATOMIC.compare_exchange(0, 1, AcqRel, Acquire).ok();
    ATOMIC.compare_exchange(0, 1, SeqCst, Acquire).ok();
    ATOMIC.compare_exchange(0, 1, SeqCst, SeqCst).ok();
    ATOMIC.compare_exchange_weak(0, 1, Relaxed, Relaxed).ok();
    ATOMIC.compare_exchange_weak(0, 1, Acquire, Relaxed).ok();
    ATOMIC.compare_exchange_weak(0, 1, Release, Relaxed).ok();
    ATOMIC.compare_exchange_weak(0, 1, AcqRel, Relaxed).ok();
    ATOMIC.compare_exchange_weak(0, 1, SeqCst, Relaxed).ok();
    ATOMIC.compare_exchange_weak(0, 1, Acquire, Acquire).ok();
    ATOMIC.compare_exchange_weak(0, 1, AcqRel, Acquire).ok();
    ATOMIC.compare_exchange_weak(0, 1, SeqCst, Acquire).ok();
    ATOMIC.compare_exchange_weak(0, 1, SeqCst, SeqCst).ok();
}
<|endoftext|>


<commit_before>#![feature(fs_walk)]
#![feature(slice_patterns)]

extern crate rustc_serialize;
extern crate time;
extern crate eve;

use std::env;
use std::fs::{OpenOptions, walk_dir};
use std::io::prelude::*;
use rustc_serialize::json::{ToJson, Json};

use eve::server::*;
use eve::flow::*;
use eve::compiler::*;

fn read_events(filename: &str) -> Vec<Event> {
    let mut events_string = String::new();
    {
        let mut events_file = OpenOptions::new().create(true).open(&filename).unwrap();
        events_file.read_to_string(&mut events_string).unwrap();
    }
    events_string.lines().map(|line|
        FromJson::from_json(&Json::from_str(&line).unwrap())
        ).collect()
}

fn write_events(filename: &str, events: &[Event]) {
    let mut new_events_file = OpenOptions::new().create(true).truncate(true).write(true).open(&filename).unwrap();
    for event in events.iter() {
        new_events_file.write_all(format!("{}", event.to_json()).as_bytes()).unwrap();
        new_events_file.write_all("\n".as_bytes()).unwrap();
    }
}

fn all_filenames() -> Vec<String> {
    let mut filenames = vec![];
    filenames.push("./events".to_owned());
    filenames.push("./boostrap".to_owned());
    for entry in walk_dir("./test-inputs").unwrap() {
        filenames.push(entry.unwrap().path().to_str().unwrap().to_owned());
    }
    for entry in walk_dir("./test-outputs").unwrap() {
        filenames.push(entry.unwrap().path().to_str().unwrap().to_owned());
    }
    for entry in walk_dir("./examples").unwrap() {
        filenames.push(entry.unwrap().path().to_str().unwrap().to_owned());
    }
    filenames
}

fn remove_view(id: &str) {
    for filename in all_filenames() {
        let mut events = read_events(&filename[..]);
        for event in events.iter_mut() {
            event.changes.retain(|&(ref change_id, _)| change_id != id);
        }
        write_events(&filename[..], &events[..]);
    }
}

fn reset_compiler() {
    let compiler_schema = compiler_schema();
    for filename in all_filenames() {
        let mut events = read_events(&filename[..]);
        for event in events.iter_mut() {
            event.changes.retain(|&(ref change_id, _)|
                !compiler_schema.iter().any(|&(ref id, _)| change_id == id)
                );
        }
        write_events(&filename[..], &events[..]);
    }
}

fn compact(filename: &str) {
    let mut flow = Flow::new();
    let bootstrap_events = read_events("./bootstrap");
    let events = read_events(&filename[..]);
    let mut flow = Flow::new();
    for event in bootstrap_events.into_iter().chain(events.into_iter()) {
        flow = flow.quiesce(event.changes);
    }
    // TODO session is blank which doesn't seem to matter because it is never used
    write_events(&filename[..], &[Event{changes: flow.as_changes(), session: "".to_owned()}]);
}

fn make_bug_test() {
    let events = read_events("./events");
    let time = time::precise_time_ns();
    let input_filename = format!("./test-inputs/bug-{}", time);
    let output_filename = format!("./test-outputs/bug-{}", time);
    write_events(&input_filename[..], &events[..]);
    write_events(&output_filename[..], &[]);
}

fn make_regression_test() {
    let mut flow = Flow::new();
    let bootstrap_events = read_events("./bootstrap");
    let events = read_events("./events");
    let mut flow = Flow::new();
    for event in bootstrap_events.into_iter().chain(events.into_iter()) {
        flow = flow.quiesce(event.changes);
    }
    let time = time::precise_time_ns();
    let input_filename = format!("./test-inputs/regression-{}", time);
    let output_filename = format!("./test-outputs/regression-{}", time);
    write_events(&input_filename[..], &[Event{changes: flow.as_changes(), session: "".to_owned()}]);
    write_events(&output_filename[..],  &[Event{changes: flow.as_changes(), session: "".to_owned()}]);
}

#[test]
fn test_examples() {
    let inputs = walk_dir("./test-inputs").unwrap().collect::<Vec<_>>();
    let outputs = walk_dir("./test-outputs").unwrap().collect::<Vec<_>>();
    assert_eq!(inputs.len(), outputs.len());
    for (input_entry, output_entry) in inputs.into_iter().zip(outputs.into_iter()) {
        let input_filename = input_entry.unwrap().path().to_str().unwrap().to_owned();
        let output_filename = output_entry.unwrap().path().to_str().unwrap().to_owned();
        println!("Testing {:?} against {:?}", input_filename, output_filename);

        let bootstrap_events = read_events("./bootstrap");
        let input_events = read_events(&input_filename[..]);
        let mut input_flow = Flow::new();
        for event in bootstrap_events.into_iter().chain(input_events.into_iter()) {
            input_flow = input_flow.quiesce(event.changes);
        }

        let output_events = read_events(&output_filename[..]);
        let mut output_flow = Flow::new();
        for event in output_events.into_iter() {
            output_flow.change(event.changes);
        }

        let schema = schema();
        for output_cell in output_flow.outputs.iter() {
            let output = output_cell.borrow();
            // ignore internal views
            if !schema.iter().any(|&(ref id, _)| **id == *output.view) {
                assert_eq!(output.index, input_flow.get_output(&output.view[..]).index);
            }
        }
    }
}

fn main() {
    let args = env::args().collect::<Vec<String>>();
    let borrowed_args = args.iter().map(|s| &s[..]).collect::<Vec<&str>>();
    match &borrowed_args[..] {
        [_, "remove_view", id] => remove_view(id),
        [_, "reset_compiler"] => reset_compiler(),
        [_, "compact", filename] => compact(filename),
        [_, "make_bug_test"] => make_bug_test(),
        [_, "make_regression_test"] => make_regression_test(),
        other => panic!("Bad arguments (look at src/bin/migrate.rs for correct usage): {:?}", &other[1..]),
    }
}<commit_msg>Reset command should also reset client state<commit_after>#![feature(fs_walk)]
#![feature(slice_patterns)]

extern crate rustc_serialize;
extern crate time;
extern crate eve;

use std::env;
use std::fs::{OpenOptions, walk_dir};
use std::io::prelude::*;
use rustc_serialize::json::{ToJson, Json};

use eve::server::*;
use eve::flow::*;
use eve::compiler::*;

fn read_events(filename: &str) -> Vec<Event> {
    let mut events_string = String::new();
    {
        let mut events_file = OpenOptions::new().create(true).open(&filename).unwrap();
        events_file.read_to_string(&mut events_string).unwrap();
    }
    events_string.lines().map(|line|
        FromJson::from_json(&Json::from_str(&line).unwrap())
        ).collect()
}

fn write_events(filename: &str, events: &[Event]) {
    let mut new_events_file = OpenOptions::new().create(true).truncate(true).write(true).open(&filename).unwrap();
    for event in events.iter() {
        new_events_file.write_all(format!("{}", event.to_json()).as_bytes()).unwrap();
        new_events_file.write_all("\n".as_bytes()).unwrap();
    }
}

fn all_filenames() -> Vec<String> {
    let mut filenames = vec![];
    filenames.push("./events".to_owned());
    filenames.push("./boostrap".to_owned());
    for entry in walk_dir("./test-inputs").unwrap() {
        filenames.push(entry.unwrap().path().to_str().unwrap().to_owned());
    }
    for entry in walk_dir("./test-outputs").unwrap() {
        filenames.push(entry.unwrap().path().to_str().unwrap().to_owned());
    }
    for entry in walk_dir("./examples").unwrap() {
        filenames.push(entry.unwrap().path().to_str().unwrap().to_owned());
    }
    filenames
}

fn remove_view(id: &str) {
    for filename in all_filenames() {
        let mut events = read_events(&filename[..]);
        for event in events.iter_mut() {
            event.changes.retain(|&(ref change_id, _)| change_id != id);
        }
        write_events(&filename[..], &events[..]);
    }
}

fn reset_internal_views() {
    let compiler_schema = compiler_schema();
    let client_schema = client_schema();
    for filename in all_filenames() {
        let mut events = read_events(&filename[..]);
        for event in events.iter_mut() {
            event.changes.retain(|&(ref change_id, _)|
                !compiler_schema.iter().chain(client_schema.iter()).any(|&(ref id, _)| change_id == id)
                );
        }
        write_events(&filename[..], &events[..]);
    }
}

fn compact(filename: &str) {
    let mut flow = Flow::new();
    let bootstrap_events = read_events("./bootstrap");
    let events = read_events(&filename[..]);
    let mut flow = Flow::new();
    for event in bootstrap_events.into_iter().chain(events.into_iter()) {
        flow = flow.quiesce(event.changes);
    }
    // TODO session is blank which doesn't seem to matter because it is never used
    write_events(&filename[..], &[Event{changes: flow.as_changes(), session: "".to_owned()}]);
}

fn make_bug_test() {
    let events = read_events("./events");
    let time = time::precise_time_ns();
    let input_filename = format!("./test-inputs/bug-{}", time);
    let output_filename = format!("./test-outputs/bug-{}", time);
    write_events(&input_filename[..], &events[..]);
    write_events(&output_filename[..], &[]);
}

fn make_regression_test() {
    let mut flow = Flow::new();
    let bootstrap_events = read_events("./bootstrap");
    let events = read_events("./events");
    let mut flow = Flow::new();
    for event in bootstrap_events.into_iter().chain(events.into_iter()) {
        flow = flow.quiesce(event.changes);
    }
    let time = time::precise_time_ns();
    let input_filename = format!("./test-inputs/regression-{}", time);
    let output_filename = format!("./test-outputs/regression-{}", time);
    write_events(&input_filename[..], &[Event{changes: flow.as_changes(), session: "".to_owned()}]);
    write_events(&output_filename[..],  &[Event{changes: flow.as_changes(), session: "".to_owned()}]);
}

#[test]
fn test_examples() {
    let inputs = walk_dir("./test-inputs").unwrap().collect::<Vec<_>>();
    let outputs = walk_dir("./test-outputs").unwrap().collect::<Vec<_>>();
    assert_eq!(inputs.len(), outputs.len());
    for (input_entry, output_entry) in inputs.into_iter().zip(outputs.into_iter()) {
        let input_filename = input_entry.unwrap().path().to_str().unwrap().to_owned();
        let output_filename = output_entry.unwrap().path().to_str().unwrap().to_owned();
        println!("Testing {:?} against {:?}", input_filename, output_filename);

        let bootstrap_events = read_events("./bootstrap");
        let input_events = read_events(&input_filename[..]);
        let mut input_flow = Flow::new();
        for event in bootstrap_events.into_iter().chain(input_events.into_iter()) {
            input_flow = input_flow.quiesce(event.changes);
        }

        let output_events = read_events(&output_filename[..]);
        let mut output_flow = Flow::new();
        for event in output_events.into_iter() {
            output_flow.change(event.changes);
        }

        let schema = schema();
        for output_cell in output_flow.outputs.iter() {
            let output = output_cell.borrow();
            // ignore internal views
            if !schema.iter().any(|&(ref id, _)| **id == *output.view) {
                assert_eq!(output.index, input_flow.get_output(&output.view[..]).index);
            }
        }
    }
}

fn main() {
    let args = env::args().collect::<Vec<String>>();
    let borrowed_args = args.iter().map(|s| &s[..]).collect::<Vec<&str>>();
    match &borrowed_args[..] {
        [_, "remove_view", id] => remove_view(id),
        [_, "reset_internal_views"] => reset_internal_views(),
        [_, "compact", filename] => compact(filename),
        [_, "make_bug_test"] => make_bug_test(),
        [_, "make_regression_test"] => make_regression_test(),
        other => panic!("Bad arguments (look at src/bin/migrate.rs for correct usage): {:?}", &other[1..]),
    }
}<|endoftext|>


<commit_before>use futures::sync::mpsc;
use futures::{try_ready, Async, AsyncSink, Future, Poll, Sink, Stream};
use std::io::{self, Read};
use tokio_postgres::types::{ToSql, Type};
use tokio_postgres::{Error, Row};
#[cfg(feature = "runtime")]
use tokio_postgres::{MakeTlsMode, Socket, TlsMode};

#[cfg(feature = "runtime")]
use crate::Builder;
use crate::{Query, Statement, Transaction};

pub struct Client(tokio_postgres::Client);

impl Client {
    #[cfg(feature = "runtime")]
    pub fn connect<T>(params: &str, tls_mode: T) -> Result<Client, Error>
    where
        T: MakeTlsMode<Socket> + 'static + Send,
        T::TlsMode: Send,
        T::Stream: Send,
        T::Future: Send,
        <T::TlsMode as TlsMode<Socket>>::Future: Send,
    {
        params.parse::<Builder>()?.connect(tls_mode)
    }

    #[cfg(feature = "runtime")]
    pub fn builder() -> Builder {
        Builder::new()
    }

    pub fn prepare(&mut self, query: &str) -> Result<Statement, Error> {
        self.0.prepare(query).wait().map(Statement)
    }

    pub fn prepare_typed(&mut self, query: &str, types: &[Type]) -> Result<Statement, Error> {
        self.0.prepare_typed(query, types).wait().map(Statement)
    }

    pub fn execute<T>(&mut self, query: &T, params: &[&dyn ToSql]) -> Result<u64, Error>
    where
        T: ?Sized + Query,
    {
        let statement = query.__statement(self)?;
        self.0.execute(&statement.0, params).wait()
    }

    pub fn query<T>(&mut self, query: &T, params: &[&dyn ToSql]) -> Result<Vec<Row>, Error>
    where
        T: ?Sized + Query,
    {
        let statement = query.__statement(self)?;
        self.0.query(&statement.0, params).collect().wait()
    }

    pub fn copy_in<T, R>(
        &mut self,
        query: &T,
        params: &[&dyn ToSql],
        reader: R,
    ) -> Result<u64, Error>
    where
        T: ?Sized + Query,
        R: Read,
    {
        let statement = query.__statement(self)?;
        let (sender, receiver) = mpsc::channel(1);
        let future = self.0.copy_in(&statement.0, params, CopyInStream(receiver));

        CopyInFuture {
            future,
            sender,
            reader,
            pending: None,
        }
        .wait()
    }

    pub fn batch_execute(&mut self, query: &str) -> Result<(), Error> {
        self.0.batch_execute(query).wait()
    }

    pub fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
        self.batch_execute("BEGIN")?;
        Ok(Transaction::new(self))
    }

    pub fn is_closed(&self) -> bool {
        self.0.is_closed()
    }
}

impl From<tokio_postgres::Client> for Client {
    fn from(c: tokio_postgres::Client) -> Client {
        Client(c)
    }
}

enum CopyData {
    Data(Vec<u8>),
    Error(io::Error),
    Done,
}

struct CopyInStream(mpsc::Receiver<CopyData>);

impl Stream for CopyInStream {
    type Item = Vec<u8>;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Vec<u8>>, io::Error> {
        match self.0.poll().expect("mpsc::Receiver can't error") {
            Async::Ready(Some(CopyData::Data(buf))) => Ok(Async::Ready(Some(buf))),
            Async::Ready(Some(CopyData::Error(e))) => Err(e),
            Async::Ready(Some(CopyData::Done)) => Ok(Async::Ready(None)),
            Async::Ready(None) => Err(io::Error::new(io::ErrorKind::Other, "writer disconnected")),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}

struct CopyInFuture<R> {
    future: tokio_postgres::CopyIn<CopyInStream>,
    sender: mpsc::Sender<CopyData>,
    reader: R,
    pending: Option<CopyData>,
}

impl<R> CopyInFuture<R> {
    fn poll_send_data(&mut self, data: CopyData) -> Poll<(), Error> {
        match self.sender.start_send(data) {
            Ok(AsyncSink::Ready) => Ok(Async::Ready(())),
            Ok(AsyncSink::NotReady(pending)) => {
                self.pending = Some(pending);
                return Ok(Async::NotReady);
            }
            // the future's hung up on its end of the channel, so we'll wait for it to report an error
            Err(_) => {
                self.pending = Some(CopyData::Done);
                return Ok(Async::NotReady);
            }
        }
    }
}

impl<R> Future for CopyInFuture<R>
where
    R: Read,
{
    type Item = u64;
    type Error = Error;

    fn poll(&mut self) -> Poll<u64, Error> {
        if let Async::Ready(n) = self.future.poll()? {
            return Ok(Async::Ready(n));
        }

        loop {
            let data = match self.pending.take() {
                Some(pending) => pending,
                None => {
                    let mut buf = vec![];
                    match self.reader.by_ref().take(4096).read_to_end(&mut buf) {
                        Ok(0) => CopyData::Done,
                        Ok(_) => CopyData::Data(buf),
                        Err(e) => CopyData::Error(e),
                    }
                }
            };

            try_ready!(self.poll_send_data(data));
        }
    }
}
<commit_msg>Fix clippy<commit_after>use futures::sync::mpsc;
use futures::{try_ready, Async, AsyncSink, Future, Poll, Sink, Stream};
use std::io::{self, Read};
use tokio_postgres::types::{ToSql, Type};
use tokio_postgres::{Error, Row};
#[cfg(feature = "runtime")]
use tokio_postgres::{MakeTlsMode, Socket, TlsMode};

#[cfg(feature = "runtime")]
use crate::Builder;
use crate::{Query, Statement, Transaction};

pub struct Client(tokio_postgres::Client);

impl Client {
    #[cfg(feature = "runtime")]
    pub fn connect<T>(params: &str, tls_mode: T) -> Result<Client, Error>
    where
        T: MakeTlsMode<Socket> + 'static + Send,
        T::TlsMode: Send,
        T::Stream: Send,
        T::Future: Send,
        <T::TlsMode as TlsMode<Socket>>::Future: Send,
    {
        params.parse::<Builder>()?.connect(tls_mode)
    }

    #[cfg(feature = "runtime")]
    pub fn builder() -> Builder {
        Builder::new()
    }

    pub fn prepare(&mut self, query: &str) -> Result<Statement, Error> {
        self.0.prepare(query).wait().map(Statement)
    }

    pub fn prepare_typed(&mut self, query: &str, types: &[Type]) -> Result<Statement, Error> {
        self.0.prepare_typed(query, types).wait().map(Statement)
    }

    pub fn execute<T>(&mut self, query: &T, params: &[&dyn ToSql]) -> Result<u64, Error>
    where
        T: ?Sized + Query,
    {
        let statement = query.__statement(self)?;
        self.0.execute(&statement.0, params).wait()
    }

    pub fn query<T>(&mut self, query: &T, params: &[&dyn ToSql]) -> Result<Vec<Row>, Error>
    where
        T: ?Sized + Query,
    {
        let statement = query.__statement(self)?;
        self.0.query(&statement.0, params).collect().wait()
    }

    pub fn copy_in<T, R>(
        &mut self,
        query: &T,
        params: &[&dyn ToSql],
        reader: R,
    ) -> Result<u64, Error>
    where
        T: ?Sized + Query,
        R: Read,
    {
        let statement = query.__statement(self)?;
        let (sender, receiver) = mpsc::channel(1);
        let future = self.0.copy_in(&statement.0, params, CopyInStream(receiver));

        CopyInFuture {
            future,
            sender,
            reader,
            pending: None,
        }
        .wait()
    }

    pub fn batch_execute(&mut self, query: &str) -> Result<(), Error> {
        self.0.batch_execute(query).wait()
    }

    pub fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
        self.batch_execute("BEGIN")?;
        Ok(Transaction::new(self))
    }

    pub fn is_closed(&self) -> bool {
        self.0.is_closed()
    }
}

impl From<tokio_postgres::Client> for Client {
    fn from(c: tokio_postgres::Client) -> Client {
        Client(c)
    }
}

enum CopyData {
    Data(Vec<u8>),
    Error(io::Error),
    Done,
}

struct CopyInStream(mpsc::Receiver<CopyData>);

impl Stream for CopyInStream {
    type Item = Vec<u8>;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Vec<u8>>, io::Error> {
        match self.0.poll().expect("mpsc::Receiver can't error") {
            Async::Ready(Some(CopyData::Data(buf))) => Ok(Async::Ready(Some(buf))),
            Async::Ready(Some(CopyData::Error(e))) => Err(e),
            Async::Ready(Some(CopyData::Done)) => Ok(Async::Ready(None)),
            Async::Ready(None) => Err(io::Error::new(io::ErrorKind::Other, "writer disconnected")),
            Async::NotReady => Ok(Async::NotReady),
        }
    }
}

struct CopyInFuture<R> {
    future: tokio_postgres::CopyIn<CopyInStream>,
    sender: mpsc::Sender<CopyData>,
    reader: R,
    pending: Option<CopyData>,
}

impl<R> CopyInFuture<R> {
    fn poll_send_data(&mut self, data: CopyData) -> Poll<(), Error> {
        match self.sender.start_send(data) {
            Ok(AsyncSink::Ready) => Ok(Async::Ready(())),
            Ok(AsyncSink::NotReady(pending)) => {
                self.pending = Some(pending);
                Ok(Async::NotReady)
            }
            // the future's hung up on its end of the channel, so we'll wait for it to report an error
            Err(_) => {
                self.pending = Some(CopyData::Done);
                Ok(Async::NotReady)
            }
        }
    }
}

impl<R> Future for CopyInFuture<R>
where
    R: Read,
{
    type Item = u64;
    type Error = Error;

    fn poll(&mut self) -> Poll<u64, Error> {
        if let Async::Ready(n) = self.future.poll()? {
            return Ok(Async::Ready(n));
        }

        loop {
            let data = match self.pending.take() {
                Some(pending) => pending,
                None => {
                    let mut buf = vec![];
                    match self.reader.by_ref().take(4096).read_to_end(&mut buf) {
                        Ok(0) => CopyData::Done,
                        Ok(_) => CopyData::Data(buf),
                        Err(e) => CopyData::Error(e),
                    }
                }
            };

            try_ready!(self.poll_send_data(data));
        }
    }
}
<|endoftext|>


<commit_before>//! Check that types of up to 128 bits are passed and returned by-value instead of via pointer.

// compile-flags: -C no-prepopulate-passes -O
// only-x86_64

#![crate_type = "lib"]

pub struct S {
    a: u64,
    b: u32,
    c: u32,
}

// CHECK: define i128 @modify(i128 %0)
#[no_mangle]
pub fn modify(s: S) -> S {
    S { a: s.a + s.a, b: s.b + s.b, c: s.c + s.c }
}

#[repr(packed)]
pub struct TooBig {
    a: u64,
    b: u32,
    c: u32,
    d: u8,
}

// CHECK: define void @m_big(%TooBig* [[ATTRS:.*sret.*]], %TooBig* [[ATTRS2:.*]] %s)
#[no_mangle]
pub fn m_big(s: TooBig) -> TooBig {
    TooBig { a: s.a + s.a, b: s.b + s.b, c: s.c + s.c, d: s.d + s.d }
}
<commit_msg>make test work in llvm 9<commit_after>//! Check that types of up to 128 bits are passed and returned by-value instead of via pointer.

// compile-flags: -C no-prepopulate-passes -O
// only-x86_64

#![crate_type = "lib"]

pub struct S {
    a: u64,
    b: u32,
    c: u32,
}

// CHECK: define i128 @modify(i128{{( %0)?}})
#[no_mangle]
pub fn modify(s: S) -> S {
    S { a: s.a + s.a, b: s.b + s.b, c: s.c + s.c }
}

#[repr(packed)]
pub struct TooBig {
    a: u64,
    b: u32,
    c: u32,
    d: u8,
}

// CHECK: define void @m_big(%TooBig* [[ATTRS:.*sret.*]], %TooBig* [[ATTRS2:.*]] %s)
#[no_mangle]
pub fn m_big(s: TooBig) -> TooBig {
    TooBig { a: s.a + s.a, b: s.b + s.b, c: s.c + s.c, d: s.d + s.d }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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
//
//     http://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 habitat_butterfly::member::Health;
use habitat_butterfly::rumor::election::ElectionStatus;

use btest;

#[test]
fn three_members_run_election() {
    let mut net = btest::SwimNet::new(3);
    net.mesh();
    net.add_service(0, "core/witcher/1.2.3/20161208121212");
    net.add_service(1, "core/witcher/1.2.3/20161208121212");
    net.add_service(2, "core/witcher/1.2.3/20161208121212");

    net.add_election(0, "witcher");
    net.add_election(1, "witcher");
    net.add_election(2, "witcher");

    assert_wait_for_election_status!(net, [0..3], "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_equal_election!(net, [0..3, 0..3], "witcher.prod");
}

#[test]
fn three_members_run_election_from_one_starting_rumor() {
    let mut net = btest::SwimNet::new(3);
    net.mesh();
    net.add_service(0, "core/witcher/1.2.3/20161208121212");
    net.add_service(1, "core/witcher/1.2.3/20161208121212");
    net.add_service(2, "core/witcher/1.2.3/20161208121212");
    net.add_election(0, "witcher");
    assert_wait_for_election_status!(net, [0..3], "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_equal_election!(net, [0..3, 0..3], "witcher.prod");
}

#[test]
fn two_members_fail_to_find_quorum() {
    let mut net = btest::SwimNet::new(2);
    net.mesh();
    net.add_service(0, "core/witcher/1.2.3/20161208121212");
    net.add_service(1, "core/witcher/1.2.3/20161208121212");
    net.add_election(0, "witcher");
    assert_wait_for_equal_election!(net, [0..2, 0..2], "witcher.prod");
    assert_wait_for_election_status!(net, [0..2], "witcher.prod", ElectionStatus::NoQuorum);
}

#[test]
fn two_members_find_quorum_when_a_third_comes() {
    let mut net = btest::SwimNet::new(2);
    net.mesh();
    net.add_service(0, "core/witcher/1.2.3/20161208121212");
    net.add_service(1, "core/witcher/1.2.3/20161208121212");
    net.add_election(0, "witcher");
    assert_wait_for_equal_election!(net, [0..2, 0..2], "witcher.prod");
    assert_wait_for_election_status!(net, [0..2], "witcher.prod", ElectionStatus::NoQuorum);

    net.members.push(btest::start_server("2", None, 0));
    net.add_service(2, "core/witcher/1.2.3/20161208121212");
    net.connect(2, 0);
    assert_wait_for_election_status!(net, [0..2], "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_equal_election!(net, [0..3, 0..3], "witcher.prod");
}

#[test]
#[ignore]
fn five_members_elect_a_new_leader_when_the_old_one_dies() {
    let mut net = btest::SwimNet::new(5);
    net.mesh();
    net.add_service(0, "core/witcher/1.2.3/20161208121212");
    net.add_service(1, "core/witcher/1.2.3/20161208121212");
    net.add_service(2, "core/witcher/1.2.3/20161208121212");
    net.add_service(3, "core/witcher/1.2.3/20161208121212");
    net.add_service(4, "core/witcher/1.2.3/20161208121212");
    net.add_election(0, "witcher");
    assert_wait_for_election_status!(net, [0..5], "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_equal_election!(net, [0..5, 0..5], "witcher.prod");

    let mut leader_id = String::from("");
    net[0]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            leader_id = e.unwrap().member_id.to_string();
        });

    let mut paused = 0;
    for (index, server) in net.iter_mut().enumerate() {
        if server.member_id() == &leader_id[..] {
            paused = index;
        }
    }
    net[paused].pause();
    let paused_id = net[paused].member_id();
    assert_wait_for_health_of!(net, paused, Health::Confirmed);
    if paused == 0 {
        net[1].restart_elections();
    } else {
        net[0].restart_elections();
    }

    for i in 0..5 {
        if !i == paused {
            assert_wait_for_election_status!(net, i, "witcher.prod", ElectionStatus::Running);
        }
    }

    for i in 0..5 {
        if !i == paused {
            assert_wait_for_election_status!(net, i, "witcher.prod", ElectionStatus::Finished);
        }
    }

    if paused == 0 {
        net[1]
            .election_store
            .with_rumor("witcher.prod", "election", |e| {
                assert!(e.unwrap().term == 1);
                assert!(e.unwrap().member_id != paused_id);
            });
    } else {
        net[0]
            .election_store
            .with_rumor("witcher.prod", "election", |e| {
                assert!(e.unwrap().term == 1);
                assert!(e.unwrap().member_id != paused_id);
            });
    }
}

#[test]
fn five_members_elect_a_new_leader_when_they_are_quorum_partitioned() {
    let mut net = btest::SwimNet::new_with_suitability(vec![1, 0, 0, 0, 0]);
    net[0]
        .member
        .write()
        .expect("Member lock is poisoned")
        .set_persistent();
    net[4]
        .member
        .write()
        .expect("Member lock is poisoned")
        .set_persistent();
    net.add_service(0, "core/witcher/1.2.3/20161208121212");
    net.add_service(1, "core/witcher/1.2.3/20161208121212");
    net.add_service(2, "core/witcher/1.2.3/20161208121212");
    net.add_service(3, "core/witcher/1.2.3/20161208121212");
    net.add_service(4, "core/witcher/1.2.3/20161208121212");
    net.add_election(0, "witcher");
    net.connect(0, 1);
    net.connect(1, 2);
    net.connect(2, 3);
    net.connect(3, 4);
    assert_wait_for_health_of!(net, [0..5, 0..5], Health::Alive);
    assert_wait_for_election_status!(net, [0..5], "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_equal_election!(net, [0..5, 0..5], "witcher.prod");

    let mut leader_id = String::from("");
    net[0]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            leader_id = e.unwrap().member_id.to_string();
        });

    assert_eq!(leader_id, net[0].member_id());

    let mut leader_index = 0;
    for (index, server) in net.iter_mut().enumerate() {
        if server.member_id() == &leader_id[..] {
            leader_index = index;
        }
    }
    println!("Leader index: {}", leader_index);

    let mut new_leader_id = String::from("");
    net.partition(0..2, 2..5);
    assert_wait_for_health_of!(net, [0..2, 2..5], Health::Confirmed);
    net[0].restart_elections();
    net[4].restart_elections();
    assert_wait_for_election_status!(net, 0, "witcher.prod", ElectionStatus::NoQuorum);
    assert_wait_for_election_status!(net, 1, "witcher.prod", ElectionStatus::NoQuorum);
    assert_wait_for_election_status!(net, 2, "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_election_status!(net, 3, "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_election_status!(net, 4, "witcher.prod", ElectionStatus::Finished);
    net[0]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            println!("OLD: {:#?}", e);
            new_leader_id = e.unwrap().member_id.to_string();
        });
    net[2]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            println!("NEW: {:#?}", e);
            new_leader_id = e.unwrap().member_id.to_string();
        });
    assert!(leader_id != new_leader_id);
    println!("Leader {} New {}", leader_id, new_leader_id);
    net.unpartition(0..2, 2..5);
    assert_wait_for_health_of!(net, [0..5, 0..5], Health::Alive);
    assert_wait_for_election_status!(net, 0, "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_election_status!(net, 1, "witcher.prod", ElectionStatus::Finished);

    net[4]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            println!("MAJORITY: {:#?}", e);
        });

    net[0]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            println!("MINORITY: {:#?}", e);
            assert_eq!(new_leader_id, e.unwrap().member_id.to_string());
        });
}
<commit_msg>Remove obsolete quorum tests<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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
//
//     http://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 habitat_butterfly::member::Health;
use habitat_butterfly::rumor::election::ElectionStatus;

use btest;

#[test]
fn three_members_run_election() {
    let mut net = btest::SwimNet::new(3);
    net.mesh();
    net.add_service(0, "core/witcher/1.2.3/20161208121212");
    net.add_service(1, "core/witcher/1.2.3/20161208121212");
    net.add_service(2, "core/witcher/1.2.3/20161208121212");

    net.add_election(0, "witcher");
    net.add_election(1, "witcher");
    net.add_election(2, "witcher");

    assert_wait_for_election_status!(net, [0..3], "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_equal_election!(net, [0..3, 0..3], "witcher.prod");
}

#[test]
fn three_members_run_election_from_one_starting_rumor() {
    let mut net = btest::SwimNet::new(3);
    net.mesh();
    net.add_service(0, "core/witcher/1.2.3/20161208121212");
    net.add_service(1, "core/witcher/1.2.3/20161208121212");
    net.add_service(2, "core/witcher/1.2.3/20161208121212");
    net.add_election(0, "witcher");
    assert_wait_for_election_status!(net, [0..3], "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_equal_election!(net, [0..3, 0..3], "witcher.prod");
}

#[test]
#[ignore]
fn five_members_elect_a_new_leader_when_the_old_one_dies() {
    let mut net = btest::SwimNet::new(5);
    net.mesh();
    net.add_service(0, "core/witcher/1.2.3/20161208121212");
    net.add_service(1, "core/witcher/1.2.3/20161208121212");
    net.add_service(2, "core/witcher/1.2.3/20161208121212");
    net.add_service(3, "core/witcher/1.2.3/20161208121212");
    net.add_service(4, "core/witcher/1.2.3/20161208121212");
    net.add_election(0, "witcher");
    assert_wait_for_election_status!(net, [0..5], "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_equal_election!(net, [0..5, 0..5], "witcher.prod");

    let mut leader_id = String::from("");
    net[0]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            leader_id = e.unwrap().member_id.to_string();
        });

    let mut paused = 0;
    for (index, server) in net.iter_mut().enumerate() {
        if server.member_id() == &leader_id[..] {
            paused = index;
        }
    }
    net[paused].pause();
    let paused_id = net[paused].member_id();
    assert_wait_for_health_of!(net, paused, Health::Confirmed);
    if paused == 0 {
        net[1].restart_elections();
    } else {
        net[0].restart_elections();
    }

    for i in 0..5 {
        if !i == paused {
            assert_wait_for_election_status!(net, i, "witcher.prod", ElectionStatus::Running);
        }
    }

    for i in 0..5 {
        if !i == paused {
            assert_wait_for_election_status!(net, i, "witcher.prod", ElectionStatus::Finished);
        }
    }

    if paused == 0 {
        net[1]
            .election_store
            .with_rumor("witcher.prod", "election", |e| {
                assert!(e.unwrap().term == 1);
                assert!(e.unwrap().member_id != paused_id);
            });
    } else {
        net[0]
            .election_store
            .with_rumor("witcher.prod", "election", |e| {
                assert!(e.unwrap().term == 1);
                assert!(e.unwrap().member_id != paused_id);
            });
    }
}

#[test]
fn five_members_elect_a_new_leader_when_they_are_quorum_partitioned() {
    let mut net = btest::SwimNet::new_with_suitability(vec![1, 0, 0, 0, 0]);
    net[0]
        .member
        .write()
        .expect("Member lock is poisoned")
        .set_persistent();
    net[4]
        .member
        .write()
        .expect("Member lock is poisoned")
        .set_persistent();
    net.add_service(0, "core/witcher/1.2.3/20161208121212");
    net.add_service(1, "core/witcher/1.2.3/20161208121212");
    net.add_service(2, "core/witcher/1.2.3/20161208121212");
    net.add_service(3, "core/witcher/1.2.3/20161208121212");
    net.add_service(4, "core/witcher/1.2.3/20161208121212");
    net.add_election(0, "witcher");
    net.connect(0, 1);
    net.connect(1, 2);
    net.connect(2, 3);
    net.connect(3, 4);
    assert_wait_for_health_of!(net, [0..5, 0..5], Health::Alive);
    assert_wait_for_election_status!(net, [0..5], "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_equal_election!(net, [0..5, 0..5], "witcher.prod");

    let mut leader_id = String::from("");
    net[0]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            leader_id = e.unwrap().member_id.to_string();
        });

    assert_eq!(leader_id, net[0].member_id());

    let mut leader_index = 0;
    for (index, server) in net.iter_mut().enumerate() {
        if server.member_id() == &leader_id[..] {
            leader_index = index;
        }
    }
    println!("Leader index: {}", leader_index);

    let mut new_leader_id = String::from("");
    net.partition(0..2, 2..5);
    assert_wait_for_health_of!(net, [0..2, 2..5], Health::Confirmed);
    net[0].restart_elections();
    net[4].restart_elections();
    assert_wait_for_election_status!(net, 0, "witcher.prod", ElectionStatus::NoQuorum);
    assert_wait_for_election_status!(net, 1, "witcher.prod", ElectionStatus::NoQuorum);
    assert_wait_for_election_status!(net, 2, "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_election_status!(net, 3, "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_election_status!(net, 4, "witcher.prod", ElectionStatus::Finished);
    net[0]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            println!("OLD: {:#?}", e);
            new_leader_id = e.unwrap().member_id.to_string();
        });
    net[2]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            println!("NEW: {:#?}", e);
            new_leader_id = e.unwrap().member_id.to_string();
        });
    assert!(leader_id != new_leader_id);
    println!("Leader {} New {}", leader_id, new_leader_id);
    net.unpartition(0..2, 2..5);
    assert_wait_for_health_of!(net, [0..5, 0..5], Health::Alive);
    assert_wait_for_election_status!(net, 0, "witcher.prod", ElectionStatus::Finished);
    assert_wait_for_election_status!(net, 1, "witcher.prod", ElectionStatus::Finished);

    net[4]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            println!("MAJORITY: {:#?}", e);
        });

    net[0]
        .election_store
        .with_rumor("witcher.prod", "election", |e| {
            println!("MINORITY: {:#?}", e);
            assert_eq!(new_leader_id, e.unwrap().member_id.to_string());
        });
}
<|endoftext|>