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

use grpcio::{ChannelBuilder, Environment};
use kvproto::{kvrpcpb::*, metapb, tikvpb::TikvClient};
use test_raftstore::*;
use tikv_util::{collections::HashMap, HandyRwLock};

use std::sync::Arc;

#[test]
fn test_physical_scan_lock() {
    let (_cluster, client, ctx) = must_new_cluster_and_kv_client();

    // Generate kvs like k10, v10, ts=10; k11, v11, ts=11; ...
    let kv: Vec<_> = (10..20)
        .map(|i| (i, vec![b'k', i as u8], vec![b'v', i as u8]))
        .collect();

    for (ts, k, v) in &kv {
        let mut mutation = Mutation::default();
        mutation.set_op(Op::Put);
        mutation.set_key(k.clone());
        mutation.set_value(v.clone());
        must_kv_prewrite(&client, ctx.clone(), vec![mutation], k.clone(), *ts);
    }

    let all_locks: Vec<_> = kv
        .into_iter()
        .map(|(ts, k, _)| {
            // Create a LockInfo that matches the prewrite request in `must_kv_prewrite`.
            let mut lock_info = LockInfo::default();
            lock_info.set_primary_lock(k.clone());
            lock_info.set_lock_version(ts);
            lock_info.set_key(k);
            lock_info.set_lock_ttl(3000);
            lock_info.set_lock_type(Op::Put);
            lock_info.set_min_commit_ts(ts + 1);
            lock_info
        })
        .collect();

    let check_result = |got_locks: &[_], expected_locks: &[_]| {
        for i in 0..std::cmp::max(got_locks.len(), expected_locks.len()) {
            assert_eq!(got_locks[i], expected_locks[i], "lock {} mismatch", i);
        }
    };

    check_result(
        &must_physical_scan_lock(&client, ctx.clone(), 30, b"", 100),
        &all_locks,
    );
    check_result(
        &must_physical_scan_lock(&client, ctx.clone(), 15, b"", 100),
        &all_locks[0..=5],
    );
    check_result(
        &must_physical_scan_lock(&client, ctx.clone(), 10, b"", 100),
        &all_locks[0..1],
    );
    check_result(
        &must_physical_scan_lock(&client, ctx.clone(), 9, b"", 100),
        &[],
    );
    check_result(
        &must_physical_scan_lock(&client, ctx, 30, &[b'k', 13], 5),
        &all_locks[3..8],
    );
}

#[test]
fn test_applied_lock_collector() {
    let mut cluster = new_server_cluster(0, 3);
    cluster.run();

    // Create all stores' clients.
    let env = Arc::new(Environment::new(1));
    let mut clients = HashMap::default();
    for node_id in cluster.get_node_ids() {
        let channel =
            ChannelBuilder::new(Arc::clone(&env)).connect(cluster.sim.rl().get_addr(node_id));
        let client = TikvClient::new(channel);
        clients.insert(node_id, client);
    }

    // Create the ctx of the first region.
    let region = cluster.get_region(b"");
    let region_id = region.get_id();
    let leader_peer = cluster.leader_of_region(region_id).unwrap();
    let leader_store_id = leader_peer.get_store_id();
    let leader_client = clients.get(&leader_store_id).unwrap();
    let mut ctx = Context::default();
    ctx.set_region_id(region_id);
    ctx.set_peer(leader_peer.clone());
    ctx.set_region_epoch(cluster.get_region_epoch(region_id));

    // It's used to make sure all stores applies all logs.
    let wait_for_apply = |cluster: &mut Cluster<_>, region: &metapb::Region| {
        let cluster = &mut *cluster;
        region.get_peers().iter().for_each(|p| {
            let resp = async_read_on_peer(cluster, p.clone(), region.clone(), b"key", true, true)
                .recv()
                .unwrap();
            assert!(!resp.get_header().has_error(), "{:?}", resp);
        });
    };

    let check_lock = |lock: &LockInfo, k: &[u8], pk: &[u8], ts| {
        assert_eq!(lock.get_key(), k);
        assert_eq!(lock.get_primary_lock(), pk);
        assert_eq!(lock.get_lock_version(), ts);
    };

    // Register lock observer at safe point 10000.
    let mut safe_point = 10000;
    clients.iter().for_each(|(_, c)| {
        // Should report error when checking non-existent observer.
        assert!(!check_lock_observer(c, safe_point).get_error().is_empty());
        must_register_lock_observer(c, safe_point);
        assert!(must_check_lock_observer(c, safe_point, true).is_empty());
    });

    // Lock observer should only collect values in lock CF.
    let key = b"key0";
    must_kv_prewrite(
        &leader_client,
        ctx.clone(),
        vec![new_mutation(Op::Put, key, &b"v".repeat(1024))],
        key.to_vec(),
        1,
    );
    must_kv_commit(&leader_client, ctx.clone(), vec![key.to_vec()], 1, 2, 2);
    wait_for_apply(&mut cluster, &region);
    clients.iter().for_each(|(_, c)| {
        let locks = must_check_lock_observer(c, safe_point, true);
        assert_eq!(locks.len(), 1);
        check_lock(&locks[0], key, key, 1);
    });

    // Lock observer shouldn't collect locks after the safe point.
    must_kv_prewrite(
        &leader_client,
        ctx.clone(),
        vec![new_mutation(Op::Put, key, b"v")],
        key.to_vec(),
        safe_point + 1,
    );
    wait_for_apply(&mut cluster, &region);
    clients.iter().for_each(|(_, c)| {
        let locks = must_check_lock_observer(c, safe_point, true);
        assert_eq!(locks.len(), 1);
        check_lock(&locks[0], key, key, 1);
    });

    // Write 999 locks whose timestamp is less than the safe point.
    let mutations = (1..1000)
        .map(|i| new_mutation(Op::Put, format!("key{}", i).as_bytes(), b"v"))
        .collect();
    must_kv_prewrite(&leader_client, ctx.clone(), mutations, b"key1".to_vec(), 10);
    wait_for_apply(&mut cluster, &region);
    clients.iter().for_each(|(_, c)| {
        let locks = must_check_lock_observer(c, safe_point, true);
        // Plus the first lock.
        assert_eq!(locks.len(), 1000);
    });

    // Add a new store and register lock observer.
    let store_id = cluster.add_new_engine();
    let channel =
        ChannelBuilder::new(Arc::clone(&env)).connect(cluster.sim.rl().get_addr(store_id));
    let client = TikvClient::new(channel);
    must_register_lock_observer(&client, safe_point);

    // Add a new peer. Lock observer should collect all locks from snapshot.
    let peer = new_peer(store_id, store_id);
    cluster.pd_client.must_add_peer(region_id, peer.clone());
    cluster.pd_client.must_none_pending_peer(peer);
    wait_for_apply(&mut cluster, &region);
    let locks = must_check_lock_observer(&client, safe_point, true);
    assert_eq!(locks.len(), 999);

    // Should be dirty when collects too many locks.
    let mutations = (1000..1100)
        .map(|i| new_mutation(Op::Put, format!("key{}", i).as_bytes(), b"v"))
        .collect();
    must_kv_prewrite(
        &leader_client,
        ctx.clone(),
        mutations,
        b"key1000".to_vec(),
        100,
    );
    wait_for_apply(&mut cluster, &region);
    clients.insert(store_id, client);
    clients.iter().for_each(|(_, c)| {
        let resp = check_lock_observer(c, safe_point);
        assert!(resp.get_error().is_empty(), "{:?}", resp.get_error());
        assert!(!resp.get_is_clean());
        // MAX_COLLECT_SIZE is 1024.
        assert_eq!(resp.get_locks().len(), 1024);
    });

    // Reregister and check. It shouldn't clean up state.
    clients.iter().for_each(|(_, c)| {
        must_register_lock_observer(c, safe_point);
        let resp = check_lock_observer(c, safe_point);
        assert!(resp.get_error().is_empty(), "{:?}", resp.get_error());
        assert!(!resp.get_is_clean());
        // MAX_COLLECT_SIZE is 1024.
        assert_eq!(resp.get_locks().len(), 1024);
    });

    // Register lock observer at a later safe point. Lock observer should reset its state.
    safe_point += 1;
    clients.iter().for_each(|(_, c)| {
        must_register_lock_observer(c, safe_point);
        assert!(must_check_lock_observer(c, safe_point, true).is_empty());
        // Can't register observer with smaller max_ts.
        assert!(!register_lock_observer(&c, safe_point - 1)
            .get_error()
            .is_empty());
        assert!(must_check_lock_observer(c, safe_point, true).is_empty());
    });
    let leader_client = clients.get(&leader_store_id).unwrap();
    must_kv_prewrite(
        &leader_client,
        ctx.clone(),
        vec![new_mutation(Op::Put, b"key1100", b"v")],
        b"key1100".to_vec(),
        safe_point,
    );
    wait_for_apply(&mut cluster, &region);
    clients.iter().for_each(|(_, c)| {
        // Should collect locks according to the new max ts.
        let locks = must_check_lock_observer(c, safe_point, true);
        assert_eq!(locks.len(), 1, "{:?}", locks);
        // Shouldn't remove it with a wrong max ts.
        assert!(!remove_lock_observer(c, safe_point - 1)
            .get_error()
            .is_empty());
        let locks = must_check_lock_observer(c, safe_point, true);
        assert_eq!(locks.len(), 1, "{:?}", locks);
        // Remove lock observers.
        must_remove_lock_observer(c, safe_point);
        assert!(!check_lock_observer(c, safe_point).get_error().is_empty());
    });
}
<commit_msg>test: stablize test_applied_lock_collector (#8369)<commit_after>// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.

use grpcio::{ChannelBuilder, Environment};
use kvproto::{kvrpcpb::*, metapb, tikvpb::TikvClient};
use test_raftstore::*;
use tikv_util::{collections::HashMap, HandyRwLock};

use std::sync::Arc;

#[test]
fn test_physical_scan_lock() {
    let (_cluster, client, ctx) = must_new_cluster_and_kv_client();

    // Generate kvs like k10, v10, ts=10; k11, v11, ts=11; ...
    let kv: Vec<_> = (10..20)
        .map(|i| (i, vec![b'k', i as u8], vec![b'v', i as u8]))
        .collect();

    for (ts, k, v) in &kv {
        let mut mutation = Mutation::default();
        mutation.set_op(Op::Put);
        mutation.set_key(k.clone());
        mutation.set_value(v.clone());
        must_kv_prewrite(&client, ctx.clone(), vec![mutation], k.clone(), *ts);
    }

    let all_locks: Vec<_> = kv
        .into_iter()
        .map(|(ts, k, _)| {
            // Create a LockInfo that matches the prewrite request in `must_kv_prewrite`.
            let mut lock_info = LockInfo::default();
            lock_info.set_primary_lock(k.clone());
            lock_info.set_lock_version(ts);
            lock_info.set_key(k);
            lock_info.set_lock_ttl(3000);
            lock_info.set_lock_type(Op::Put);
            lock_info.set_min_commit_ts(ts + 1);
            lock_info
        })
        .collect();

    let check_result = |got_locks: &[_], expected_locks: &[_]| {
        for i in 0..std::cmp::max(got_locks.len(), expected_locks.len()) {
            assert_eq!(got_locks[i], expected_locks[i], "lock {} mismatch", i);
        }
    };

    check_result(
        &must_physical_scan_lock(&client, ctx.clone(), 30, b"", 100),
        &all_locks,
    );
    check_result(
        &must_physical_scan_lock(&client, ctx.clone(), 15, b"", 100),
        &all_locks[0..=5],
    );
    check_result(
        &must_physical_scan_lock(&client, ctx.clone(), 10, b"", 100),
        &all_locks[0..1],
    );
    check_result(
        &must_physical_scan_lock(&client, ctx.clone(), 9, b"", 100),
        &[],
    );
    check_result(
        &must_physical_scan_lock(&client, ctx, 30, &[b'k', 13], 5),
        &all_locks[3..8],
    );
}

#[test]
fn test_applied_lock_collector() {
    let mut cluster = new_server_cluster(0, 3);
    cluster.pd_client.disable_default_operator();
    cluster.run();

    // Create all stores' clients.
    let env = Arc::new(Environment::new(1));
    let mut clients = HashMap::default();
    for node_id in cluster.get_node_ids() {
        let channel =
            ChannelBuilder::new(Arc::clone(&env)).connect(cluster.sim.rl().get_addr(node_id));
        let client = TikvClient::new(channel);
        clients.insert(node_id, client);
    }

    // Create the ctx of the first region.
    let region = cluster.get_region(b"");
    let region_id = region.get_id();
    let leader_peer = cluster.leader_of_region(region_id).unwrap();
    let leader_store_id = leader_peer.get_store_id();
    let leader_client = clients.get(&leader_store_id).unwrap();
    let mut ctx = Context::default();
    ctx.set_region_id(region_id);
    ctx.set_peer(leader_peer.clone());
    ctx.set_region_epoch(cluster.get_region_epoch(region_id));

    // It's used to make sure all stores applies all logs.
    let wait_for_apply = |cluster: &mut Cluster<_>, region: &metapb::Region| {
        let cluster = &mut *cluster;
        region.get_peers().iter().for_each(|p| {
            let resp = async_read_on_peer(cluster, p.clone(), region.clone(), b"key", true, true)
                .recv()
                .unwrap();
            assert!(!resp.get_header().has_error(), "{:?}", resp);
        });
    };

    let check_lock = |lock: &LockInfo, k: &[u8], pk: &[u8], ts| {
        assert_eq!(lock.get_key(), k);
        assert_eq!(lock.get_primary_lock(), pk);
        assert_eq!(lock.get_lock_version(), ts);
    };

    // Register lock observer at safe point 10000.
    let mut safe_point = 10000;
    clients.iter().for_each(|(_, c)| {
        // Should report error when checking non-existent observer.
        assert!(!check_lock_observer(c, safe_point).get_error().is_empty());
        must_register_lock_observer(c, safe_point);
        assert!(must_check_lock_observer(c, safe_point, true).is_empty());
    });

    // Lock observer should only collect values in lock CF.
    let key = b"key0";
    must_kv_prewrite(
        &leader_client,
        ctx.clone(),
        vec![new_mutation(Op::Put, key, &b"v".repeat(1024))],
        key.to_vec(),
        1,
    );
    must_kv_commit(&leader_client, ctx.clone(), vec![key.to_vec()], 1, 2, 2);
    wait_for_apply(&mut cluster, &region);
    clients.iter().for_each(|(_, c)| {
        let locks = must_check_lock_observer(c, safe_point, true);
        assert_eq!(locks.len(), 1);
        check_lock(&locks[0], key, key, 1);
    });

    // Lock observer shouldn't collect locks after the safe point.
    must_kv_prewrite(
        &leader_client,
        ctx.clone(),
        vec![new_mutation(Op::Put, key, b"v")],
        key.to_vec(),
        safe_point + 1,
    );
    wait_for_apply(&mut cluster, &region);
    clients.iter().for_each(|(_, c)| {
        let locks = must_check_lock_observer(c, safe_point, true);
        assert_eq!(locks.len(), 1);
        check_lock(&locks[0], key, key, 1);
    });

    // Write 999 locks whose timestamp is less than the safe point.
    let mutations = (1..1000)
        .map(|i| new_mutation(Op::Put, format!("key{}", i).as_bytes(), b"v"))
        .collect();
    must_kv_prewrite(&leader_client, ctx.clone(), mutations, b"key1".to_vec(), 10);
    wait_for_apply(&mut cluster, &region);
    clients.iter().for_each(|(_, c)| {
        let locks = must_check_lock_observer(c, safe_point, true);
        // Plus the first lock.
        assert_eq!(locks.len(), 1000);
    });

    // Add a new store and register lock observer.
    let store_id = cluster.add_new_engine();
    let channel =
        ChannelBuilder::new(Arc::clone(&env)).connect(cluster.sim.rl().get_addr(store_id));
    let client = TikvClient::new(channel);
    must_register_lock_observer(&client, safe_point);

    // Add a new peer. Lock observer should collect all locks from snapshot.
    let peer = new_peer(store_id, store_id);
    cluster.pd_client.must_add_peer(region_id, peer.clone());
    cluster.pd_client.must_none_pending_peer(peer);
    wait_for_apply(&mut cluster, &region);
    let locks = must_check_lock_observer(&client, safe_point, true);
    assert_eq!(locks.len(), 999);

    // Should be dirty when collects too many locks.
    let mutations = (1000..1100)
        .map(|i| new_mutation(Op::Put, format!("key{}", i).as_bytes(), b"v"))
        .collect();
    must_kv_prewrite(
        &leader_client,
        ctx.clone(),
        mutations,
        b"key1000".to_vec(),
        100,
    );
    wait_for_apply(&mut cluster, &region);
    clients.insert(store_id, client);
    clients.iter().for_each(|(_, c)| {
        let resp = check_lock_observer(c, safe_point);
        assert!(resp.get_error().is_empty(), "{:?}", resp.get_error());
        assert!(!resp.get_is_clean());
        // MAX_COLLECT_SIZE is 1024.
        assert_eq!(resp.get_locks().len(), 1024);
    });

    // Reregister and check. It shouldn't clean up state.
    clients.iter().for_each(|(_, c)| {
        must_register_lock_observer(c, safe_point);
        let resp = check_lock_observer(c, safe_point);
        assert!(resp.get_error().is_empty(), "{:?}", resp.get_error());
        assert!(!resp.get_is_clean());
        // MAX_COLLECT_SIZE is 1024.
        assert_eq!(resp.get_locks().len(), 1024);
    });

    // Register lock observer at a later safe point. Lock observer should reset its state.
    safe_point += 1;
    clients.iter().for_each(|(_, c)| {
        must_register_lock_observer(c, safe_point);
        assert!(must_check_lock_observer(c, safe_point, true).is_empty());
        // Can't register observer with smaller max_ts.
        assert!(!register_lock_observer(&c, safe_point - 1)
            .get_error()
            .is_empty());
        assert!(must_check_lock_observer(c, safe_point, true).is_empty());
    });
    let leader_client = clients.get(&leader_store_id).unwrap();
    must_kv_prewrite(
        &leader_client,
        ctx.clone(),
        vec![new_mutation(Op::Put, b"key1100", b"v")],
        b"key1100".to_vec(),
        safe_point,
    );
    wait_for_apply(&mut cluster, &region);
    clients.iter().for_each(|(_, c)| {
        // Should collect locks according to the new max ts.
        let locks = must_check_lock_observer(c, safe_point, true);
        assert_eq!(locks.len(), 1, "{:?}", locks);
        // Shouldn't remove it with a wrong max ts.
        assert!(!remove_lock_observer(c, safe_point - 1)
            .get_error()
            .is_empty());
        let locks = must_check_lock_observer(c, safe_point, true);
        assert_eq!(locks.len(), 1, "{:?}", locks);
        // Remove lock observers.
        must_remove_lock_observer(c, safe_point);
        assert!(!check_lock_observer(c, safe_point).get_error().is_empty());
    });
}
<|endoftext|>


<commit_before>//#[macro_use]
//use adapton_syntax::* ;
use adapton_sigs::* ;
// use adapton_state::AdaptonState ;

trait List {
    type Hd ;
    type List ;
    fn nil <A:Adapton> (&mut A) -> Self::List ;
    fn cons<A:Adapton> (&mut A, Self::Hd, Self::List) -> Self::List ;

    // requisite "adaptonic" constructors: `name` and `art`:
    fn name<A:Adapton> (&mut A, A::Name, Self::List) -> Self::List ;
    fn art<A:Adapton>  (&mut A, Art<Self::List,A::Loc>) -> Self::List ;
}

trait Tree {
    type Leaf ;
    type Bin ;
    type Tree ;
    fn nil <A:Adapton> (&mut A) -> Self::Tree ;
    fn leaf<A:Adapton> (&mut A, Self::Leaf) -> Self::Tree ;
    fn bin<A:Adapton>  (&mut A, Self::Bin, Self::Tree, Self::Tree) -> Self::Tree ;

    // requisite "adaptonic" constructors: `name` and `art`:
    fn name<A:Adapton> (&mut A, A::Name, Self::Tree, Self::Tree) -> Self::Tree ;
    fn art<A:Adapton>  (&mut A, Art<Self::Tree,A::Loc>) -> Self::Tree ;
}
<commit_msg>wrote types for list and tree folds<commit_after>//#[macro_use]
//use adapton_syntax::* ;
use adapton_sigs::* ;
// use adapton_state::AdaptonState ;

trait List {
    type Hd ;
    type List ;
    fn nil <A:Adapton> (&mut A) -> Self::List ;
    fn cons<A:Adapton> (&mut A, Self::Hd, Self::List) -> Self::List ;

    // requisite "adaptonic" constructors: `name` and `art`:
    fn name<A:Adapton> (&mut A, A::Name, Self::List) -> Self::List ;
    fn art<A:Adapton>  (&mut A, Art<Self::List,A::Loc>) -> Self::List ;

    fn fold<A:Adapton,Res,Cons> (&mut A, Self::List, Res, Cons) -> Res
        where Cons:Fn(&mut A, Res, Self::Hd) -> Res ;
}

trait Tree {
    type Leaf ;
    type Bin ;
    type Tree ;
    fn nil <A:Adapton> (&mut A) -> Self::Tree ;
    fn leaf<A:Adapton> (&mut A, Self::Leaf) -> Self::Tree ;
    fn bin<A:Adapton>  (&mut A, Self::Bin, Self::Tree, Self::Tree) -> Self::Tree ;

    // requisite "adaptonic" constructors: `name` and `art`:
    fn name<A:Adapton> (&mut A, A::Name, Self::Tree, Self::Tree) -> Self::Tree ;
    fn art<A:Adapton>  (&mut A, Art<Self::Tree,A::Loc>) -> Self::Tree ;

    fn fold<A:Adapton,Res,Leaf,Bin> (&mut A, Self::Tree, Res, Leaf, Bin) -> Res
        where Leaf:Fn(&mut A, Res, Self::Leaf) -> Res
        ,      Bin:Fn(&mut A, Res, Self::Bin ) -> Res ;

    fn fold2<A:Adapton,Arg,Res,Leaf,Bin> (&mut A, Self::Tree, Arg, Leaf, Bin) -> Res
        where Leaf:Fn(&mut A, Arg, Self::Leaf) -> Res
        ,      Bin:Fn(&mut A, Arg, Self::Bin, Res, Res ) -> Res ;
}
<|endoftext|>


<commit_before>/*
 * Copyright 2022 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     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 std::{pin::Pin, sync::Arc};

use cached::Cached;
use futures::Stream;
use tokio_stream::StreamExt;
use tracing_futures::Instrument;

use crate::{
    config::Config,
    xds::{
        metrics,
        service::discovery::v3::{
            aggregated_discovery_service_server::{
                AggregatedDiscoveryService, AggregatedDiscoveryServiceServer,
            },
            DeltaDiscoveryRequest, DeltaDiscoveryResponse, DiscoveryRequest, DiscoveryResponse,
        },
        ResourceType,
    },
};

#[tracing::instrument(skip_all)]
pub async fn spawn(config: std::sync::Arc<crate::Config>) -> crate::Result<()> {
    let port = config.proxy.load().port;

    let server = AggregatedDiscoveryServiceServer::new(ControlPlane::from_arc(config));
    let server = tonic::transport::Server::builder().add_service(server);
    tracing::info!("Serving management server at {}", port);
    Ok(server
        .serve((std::net::Ipv4Addr::UNSPECIFIED, port).into())
        .await?)
}

#[derive(Clone)]
pub struct ControlPlane {
    config: Arc<Config>,
    watchers: Arc<crate::xds::resource::ResourceMap<Watchers>>,
}

struct Watchers {
    sender: tokio::sync::watch::Sender<()>,
    version: std::sync::atomic::AtomicU64,
}

impl Default for Watchers {
    fn default() -> Self {
        Self {
            sender: tokio::sync::watch::channel(()).0,
            version: <_>::default(),
        }
    }
}

impl ControlPlane {
    /// Creates a new server for managing [`Config`].
    pub fn new(config: Config) -> Self {
        Self::from_arc(Arc::new(config))
    }

    pub fn from_arc(config: Arc<Config>) -> Self {
        let this = Self {
            config,
            watchers: <_>::default(),
        };

        this.config.clusters.watch({
            let this = this.clone();
            move |_| {
                this.push_update(ResourceType::Endpoint);
                this.push_update(ResourceType::Cluster);
            }
        });

        this.config.filters.watch({
            let this = this.clone();
            move |_| {
                this.push_update(ResourceType::Listener);
            }
        });

        this
    }

    fn push_update(&self, resource_type: ResourceType) {
        let watchers = &self.watchers[resource_type];
        watchers
            .version
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        tracing::trace!(%resource_type, watchers=watchers.sender.receiver_count(), "pushing update");
        let _ = watchers.sender.send(());
    }

    fn discovery_response(
        &self,
        id: &str,
        resource_type: ResourceType,
        names: &[String],
    ) -> Result<DiscoveryResponse, tonic::Status> {
        let mut response = self
            .config
            .discovery_request(id, resource_type, names)
            .map_err(|error| tonic::Status::internal(error.to_string()))?;
        let watchers = &self.watchers[resource_type];

        let nonce = uuid::Uuid::new_v4();
        response.version_info = watchers
            .version
            .load(std::sync::atomic::Ordering::Relaxed)
            .to_string();
        response.control_plane = Some(crate::xds::config::core::v3::ControlPlane {
            identifier: self.config.proxy.load().id.clone(),
        });
        response.nonce = nonce.to_string();

        tracing::trace!(
            id = &*response.version_info,
            r#type = &*response.type_url,
            nonce = &*response.nonce,
            "discovery response"
        );

        Ok(response)
    }

    pub async fn stream_aggregated_resources(
        &self,
        mut streaming: Pin<Box<dyn Stream<Item = Result<DiscoveryRequest, tonic::Status>> + Send>>,
    ) -> Result<
        Pin<Box<dyn Stream<Item = Result<DiscoveryResponse, tonic::Status>> + Send>>,
        tonic::Status,
    > {
        tracing::trace!("starting stream");
        let message = streaming.next().await.ok_or_else(|| {
            tracing::error!("No message found");
            tonic::Status::invalid_argument("No message found")
        })??;

        if message.node.is_none() {
            tracing::error!("Node identifier was not found");
            return Err(tonic::Status::invalid_argument("Node identifier required"));
        }

        let node = message.node.clone().unwrap();
        let resource_type: ResourceType = message.type_url.parse()?;
        tracing::trace!(id = %node.id, %resource_type, "initial request");
        metrics::DISCOVERY_REQUESTS
            .with_label_values(&[&*node.id, resource_type.type_url()])
            .inc();
        let mut rx = self.watchers[resource_type].sender.subscribe();
        let mut pending_acks = cached::TimedCache::with_lifespan(1);
        let this = Self::clone(self);
        let response = this.discovery_response(&node.id, resource_type, &message.resource_names)?;
        pending_acks.cache_set(response.nonce.clone(), ());

        let id = node.id.clone();
        Ok(Box::pin(async_stream::try_stream! {
            yield response;

            loop {
                tokio::select! {
                    _ = rx.changed() => {
                        yield this.discovery_response(&id, resource_type, &message.resource_names).map(|response| {
                            pending_acks.cache_set(response.nonce.clone(), ());
                            response
                        })?;
                    }
                    new_message = streaming.next() => {
                        let new_message = match new_message.transpose() {
                            Ok(Some(value)) => value,
                            Ok(None) => break,
                            Err(error) => {
                                tracing::error!(%error, "unknown resource type");
                                continue;
                            }
                        };

                        let id = new_message.node.as_ref().map(|node| &*node.id).unwrap_or(&*id);
                        let resource_type = match new_message.type_url.parse::<ResourceType>() {
                            Ok(value) => value,
                            Err(error) => {
                                tracing::error!(%error, "unknown resource type");
                                continue;
                            }
                        };

                        tracing::trace!("new request");
                        metrics::DISCOVERY_REQUESTS.with_label_values(&[id, resource_type.type_url()]).inc();

                        if let Some(error) = &new_message.error_detail {
                            metrics::NACKS.with_label_values(&[id, resource_type.type_url()]).inc();
                            tracing::error!(nonce = %new_message.response_nonce, ?error, "NACK");
                            // Currently just resend previous discovery response.
                        } else if uuid::Uuid::parse_str(&new_message.response_nonce).is_ok() {
                            if pending_acks.cache_get(&new_message.response_nonce).is_some() {
                                tracing::info!(nonce = %new_message.response_nonce, "ACK");
                                continue
                            } else {
                                tracing::trace!(nonce = %new_message.response_nonce, "Unknown nonce: could not be found in cache");
                                continue
                            }
                        }

                        yield this.discovery_response(id, resource_type, &message.resource_names).map(|response| {
                            pending_acks.cache_set(response.nonce.clone(), ());
                            response
                        }).unwrap();
                    }
                }
            }

            tracing::info!("terminating stream");
        }.instrument(tracing::info_span!("xds_stream", %node.id, %resource_type))))
    }
}

#[tonic::async_trait]
impl AggregatedDiscoveryService for ControlPlane {
    type StreamAggregatedResourcesStream =
        std::pin::Pin<Box<dyn Stream<Item = Result<DiscoveryResponse, tonic::Status>> + Send>>;
    type DeltaAggregatedResourcesStream =
        tokio_stream::wrappers::ReceiverStream<Result<DeltaDiscoveryResponse, tonic::Status>>;

    #[tracing::instrument(skip_all)]
    async fn stream_aggregated_resources(
        &self,
        request: tonic::Request<tonic::Streaming<DiscoveryRequest>>,
    ) -> Result<tonic::Response<Self::StreamAggregatedResourcesStream>, tonic::Status> {
        Ok(tonic::Response::new(
            self.stream_aggregated_resources(Box::pin(request.into_inner()))
                .in_current_span()
                .await?,
        ))
    }

    async fn delta_aggregated_resources(
        &self,
        _request: tonic::Request<tonic::Streaming<DeltaDiscoveryRequest>>,
    ) -> Result<tonic::Response<Self::DeltaAggregatedResourcesStream>, tonic::Status> {
        Err(tonic::Status::unimplemented(
            "Quilkin doesn't currently support Delta xDS",
        ))
    }
}

#[cfg(test)]
mod tests {
    use pretty_assertions::assert_eq;
    use tokio::time::timeout;

    use super::*;
    use crate::xds::{
        config::{
            core::v3::Node,
            listener::v3::{FilterChain, Listener},
        },
        service::discovery::v3::DiscoveryResponse,
        ResourceType,
    };

    const TIMEOUT_DURATION: std::time::Duration = std::time::Duration::from_secs(10);

    #[tokio::test]
    async fn valid_response() {
        const RESOURCE: ResourceType = ResourceType::Endpoint;
        const LISTENER_TYPE: ResourceType = ResourceType::Listener;

        let mut response = DiscoveryResponse {
            version_info: String::new(),
            resources: vec![],
            type_url: RESOURCE.type_url().into(),
            ..<_>::default()
        };

        let mut listener_response = DiscoveryResponse {
            version_info: String::new(),
            resources: vec![prost_types::Any {
                type_url: LISTENER_TYPE.type_url().into(),
                value: crate::prost::encode(&Listener {
                    filter_chains: vec![FilterChain {
                        filters: vec![],
                        ..<_>::default()
                    }],
                    ..<_>::default()
                })
                .unwrap(),
            }],
            type_url: LISTENER_TYPE.type_url().into(),
            ..<_>::default()
        };

        let config = Arc::new(Config::default());
        let client = ControlPlane::from_arc(config.clone());
        let (tx, rx) = tokio::sync::mpsc::channel(256);

        let mut request = DiscoveryRequest {
            node: Some(Node {
                id: "quilkin".into(),
                user_agent_name: "quilkin".into(),
                ..Node::default()
            }),
            resource_names: vec![],
            type_url: RESOURCE.type_url().into(),
            ..DiscoveryRequest::default()
        };

        let mut listener_request = DiscoveryRequest {
            node: Some(Node {
                id: "quilkin".into(),
                user_agent_name: "quilkin".into(),
                ..Node::default()
            }),
            resource_names: vec![],
            type_url: LISTENER_TYPE.type_url().into(),
            ..DiscoveryRequest::default()
        };

        timeout(TIMEOUT_DURATION, tx.send(Ok(request.clone())))
            .await
            .unwrap()
            .unwrap();

        let mut stream = timeout(
            TIMEOUT_DURATION,
            client.stream_aggregated_resources(Box::pin(
                tokio_stream::wrappers::ReceiverStream::new(rx),
            )),
        )
        .await
        .unwrap()
        .unwrap();

        let message = timeout(TIMEOUT_DURATION, stream.next())
            .await
            .unwrap()
            .unwrap()
            .unwrap();
        response.version_info = message.version_info.clone();
        response.nonce = message.nonce.clone();
        response.control_plane = message.control_plane.clone();
        request.response_nonce = message.nonce.clone();

        assert_eq!(response, message);

        timeout(TIMEOUT_DURATION, tx.send(Ok(request.clone())))
            .await
            .unwrap()
            .unwrap();

        timeout(TIMEOUT_DURATION, tx.send(Ok(listener_request.clone())))
            .await
            .unwrap()
            .unwrap();

        let message = timeout(TIMEOUT_DURATION, stream.next())
            .await
            .unwrap()
            .unwrap()
            .unwrap();
        listener_response.control_plane = message.control_plane.clone();
        listener_response.version_info = message.version_info.clone();
        listener_response.nonce = message.nonce.clone();
        listener_request.response_nonce = message.nonce.clone();

        assert_eq!(listener_response, message);

        timeout(TIMEOUT_DURATION, tx.send(Ok(listener_request.clone())))
            .await
            .unwrap()
            .unwrap();
    }
}
<commit_msg>Use `TimedSizedCache` (#616)<commit_after>/*
 * Copyright 2022 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     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 std::{pin::Pin, sync::Arc};

use cached::Cached;
use futures::Stream;
use tokio_stream::StreamExt;
use tracing_futures::Instrument;

use crate::{
    config::Config,
    xds::{
        metrics,
        service::discovery::v3::{
            aggregated_discovery_service_server::{
                AggregatedDiscoveryService, AggregatedDiscoveryServiceServer,
            },
            DeltaDiscoveryRequest, DeltaDiscoveryResponse, DiscoveryRequest, DiscoveryResponse,
        },
        ResourceType,
    },
};

#[tracing::instrument(skip_all)]
pub async fn spawn(config: std::sync::Arc<crate::Config>) -> crate::Result<()> {
    let port = config.proxy.load().port;

    let server = AggregatedDiscoveryServiceServer::new(ControlPlane::from_arc(config));
    let server = tonic::transport::Server::builder().add_service(server);
    tracing::info!("Serving management server at {}", port);
    Ok(server
        .serve((std::net::Ipv4Addr::UNSPECIFIED, port).into())
        .await?)
}

#[derive(Clone)]
pub struct ControlPlane {
    config: Arc<Config>,
    watchers: Arc<crate::xds::resource::ResourceMap<Watchers>>,
}

struct Watchers {
    sender: tokio::sync::watch::Sender<()>,
    version: std::sync::atomic::AtomicU64,
}

impl Default for Watchers {
    fn default() -> Self {
        Self {
            sender: tokio::sync::watch::channel(()).0,
            version: <_>::default(),
        }
    }
}

impl ControlPlane {
    /// Creates a new server for managing [`Config`].
    pub fn new(config: Config) -> Self {
        Self::from_arc(Arc::new(config))
    }

    pub fn from_arc(config: Arc<Config>) -> Self {
        let this = Self {
            config,
            watchers: <_>::default(),
        };

        this.config.clusters.watch({
            let this = this.clone();
            move |_| {
                this.push_update(ResourceType::Endpoint);
                this.push_update(ResourceType::Cluster);
            }
        });

        this.config.filters.watch({
            let this = this.clone();
            move |_| {
                this.push_update(ResourceType::Listener);
            }
        });

        this
    }

    fn push_update(&self, resource_type: ResourceType) {
        let watchers = &self.watchers[resource_type];
        watchers
            .version
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        tracing::trace!(%resource_type, watchers=watchers.sender.receiver_count(), "pushing update");
        let _ = watchers.sender.send(());
    }

    fn discovery_response(
        &self,
        id: &str,
        resource_type: ResourceType,
        names: &[String],
    ) -> Result<DiscoveryResponse, tonic::Status> {
        let mut response = self
            .config
            .discovery_request(id, resource_type, names)
            .map_err(|error| tonic::Status::internal(error.to_string()))?;
        let watchers = &self.watchers[resource_type];

        let nonce = uuid::Uuid::new_v4();
        response.version_info = watchers
            .version
            .load(std::sync::atomic::Ordering::Relaxed)
            .to_string();
        response.control_plane = Some(crate::xds::config::core::v3::ControlPlane {
            identifier: self.config.proxy.load().id.clone(),
        });
        response.nonce = nonce.to_string();

        tracing::trace!(
            id = &*response.version_info,
            r#type = &*response.type_url,
            nonce = &*response.nonce,
            "discovery response"
        );

        Ok(response)
    }

    pub async fn stream_aggregated_resources(
        &self,
        mut streaming: Pin<Box<dyn Stream<Item = Result<DiscoveryRequest, tonic::Status>> + Send>>,
    ) -> Result<
        Pin<Box<dyn Stream<Item = Result<DiscoveryResponse, tonic::Status>> + Send>>,
        tonic::Status,
    > {
        tracing::trace!("starting stream");
        let message = streaming.next().await.ok_or_else(|| {
            tracing::error!("No message found");
            tonic::Status::invalid_argument("No message found")
        })??;

        if message.node.is_none() {
            tracing::error!("Node identifier was not found");
            return Err(tonic::Status::invalid_argument("Node identifier required"));
        }

        let node = message.node.clone().unwrap();
        let resource_type: ResourceType = message.type_url.parse()?;
        tracing::trace!(id = %node.id, %resource_type, "initial request");
        metrics::DISCOVERY_REQUESTS
            .with_label_values(&[&*node.id, resource_type.type_url()])
            .inc();
        let mut rx = self.watchers[resource_type].sender.subscribe();
        let mut pending_acks = cached::TimedSizedCache::with_size_and_lifespan(50, 1);
        let this = Self::clone(self);
        let response = this.discovery_response(&node.id, resource_type, &message.resource_names)?;
        pending_acks.cache_set(response.nonce.clone(), ());

        let id = node.id.clone();
        Ok(Box::pin(async_stream::try_stream! {
            yield response;

            loop {
                tokio::select! {
                    _ = rx.changed() => {
                        yield this.discovery_response(&id, resource_type, &message.resource_names).map(|response| {
                            pending_acks.cache_set(response.nonce.clone(), ());
                            response
                        })?;
                    }
                    new_message = streaming.next() => {
                        let new_message = match new_message.transpose() {
                            Ok(Some(value)) => value,
                            Ok(None) => break,
                            Err(error) => {
                                tracing::error!(%error, "unknown resource type");
                                continue;
                            }
                        };

                        let id = new_message.node.as_ref().map(|node| &*node.id).unwrap_or(&*id);
                        let resource_type = match new_message.type_url.parse::<ResourceType>() {
                            Ok(value) => value,
                            Err(error) => {
                                tracing::error!(%error, "unknown resource type");
                                continue;
                            }
                        };

                        tracing::trace!("new request");
                        metrics::DISCOVERY_REQUESTS.with_label_values(&[id, resource_type.type_url()]).inc();

                        if let Some(error) = &new_message.error_detail {
                            metrics::NACKS.with_label_values(&[id, resource_type.type_url()]).inc();
                            tracing::error!(nonce = %new_message.response_nonce, ?error, "NACK");
                            // Currently just resend previous discovery response.
                        } else if uuid::Uuid::parse_str(&new_message.response_nonce).is_ok() {
                            if pending_acks.cache_get(&new_message.response_nonce).is_some() {
                                tracing::info!(nonce = %new_message.response_nonce, "ACK");
                                continue
                            } else {
                                tracing::trace!(nonce = %new_message.response_nonce, "Unknown nonce: could not be found in cache");
                                continue
                            }
                        }

                        yield this.discovery_response(id, resource_type, &message.resource_names).map(|response| {
                            pending_acks.cache_set(response.nonce.clone(), ());
                            response
                        }).unwrap();
                    }
                }
            }

            tracing::info!("terminating stream");
        }.instrument(tracing::info_span!("xds_stream", %node.id, %resource_type))))
    }
}

#[tonic::async_trait]
impl AggregatedDiscoveryService for ControlPlane {
    type StreamAggregatedResourcesStream =
        std::pin::Pin<Box<dyn Stream<Item = Result<DiscoveryResponse, tonic::Status>> + Send>>;
    type DeltaAggregatedResourcesStream =
        tokio_stream::wrappers::ReceiverStream<Result<DeltaDiscoveryResponse, tonic::Status>>;

    #[tracing::instrument(skip_all)]
    async fn stream_aggregated_resources(
        &self,
        request: tonic::Request<tonic::Streaming<DiscoveryRequest>>,
    ) -> Result<tonic::Response<Self::StreamAggregatedResourcesStream>, tonic::Status> {
        Ok(tonic::Response::new(
            self.stream_aggregated_resources(Box::pin(request.into_inner()))
                .in_current_span()
                .await?,
        ))
    }

    async fn delta_aggregated_resources(
        &self,
        _request: tonic::Request<tonic::Streaming<DeltaDiscoveryRequest>>,
    ) -> Result<tonic::Response<Self::DeltaAggregatedResourcesStream>, tonic::Status> {
        Err(tonic::Status::unimplemented(
            "Quilkin doesn't currently support Delta xDS",
        ))
    }
}

#[cfg(test)]
mod tests {
    use pretty_assertions::assert_eq;
    use tokio::time::timeout;

    use super::*;
    use crate::xds::{
        config::{
            core::v3::Node,
            listener::v3::{FilterChain, Listener},
        },
        service::discovery::v3::DiscoveryResponse,
        ResourceType,
    };

    const TIMEOUT_DURATION: std::time::Duration = std::time::Duration::from_secs(10);

    #[tokio::test]
    async fn valid_response() {
        const RESOURCE: ResourceType = ResourceType::Endpoint;
        const LISTENER_TYPE: ResourceType = ResourceType::Listener;

        let mut response = DiscoveryResponse {
            version_info: String::new(),
            resources: vec![],
            type_url: RESOURCE.type_url().into(),
            ..<_>::default()
        };

        let mut listener_response = DiscoveryResponse {
            version_info: String::new(),
            resources: vec![prost_types::Any {
                type_url: LISTENER_TYPE.type_url().into(),
                value: crate::prost::encode(&Listener {
                    filter_chains: vec![FilterChain {
                        filters: vec![],
                        ..<_>::default()
                    }],
                    ..<_>::default()
                })
                .unwrap(),
            }],
            type_url: LISTENER_TYPE.type_url().into(),
            ..<_>::default()
        };

        let config = Arc::new(Config::default());
        let client = ControlPlane::from_arc(config.clone());
        let (tx, rx) = tokio::sync::mpsc::channel(256);

        let mut request = DiscoveryRequest {
            node: Some(Node {
                id: "quilkin".into(),
                user_agent_name: "quilkin".into(),
                ..Node::default()
            }),
            resource_names: vec![],
            type_url: RESOURCE.type_url().into(),
            ..DiscoveryRequest::default()
        };

        let mut listener_request = DiscoveryRequest {
            node: Some(Node {
                id: "quilkin".into(),
                user_agent_name: "quilkin".into(),
                ..Node::default()
            }),
            resource_names: vec![],
            type_url: LISTENER_TYPE.type_url().into(),
            ..DiscoveryRequest::default()
        };

        timeout(TIMEOUT_DURATION, tx.send(Ok(request.clone())))
            .await
            .unwrap()
            .unwrap();

        let mut stream = timeout(
            TIMEOUT_DURATION,
            client.stream_aggregated_resources(Box::pin(
                tokio_stream::wrappers::ReceiverStream::new(rx),
            )),
        )
        .await
        .unwrap()
        .unwrap();

        let message = timeout(TIMEOUT_DURATION, stream.next())
            .await
            .unwrap()
            .unwrap()
            .unwrap();
        response.version_info = message.version_info.clone();
        response.nonce = message.nonce.clone();
        response.control_plane = message.control_plane.clone();
        request.response_nonce = message.nonce.clone();

        assert_eq!(response, message);

        timeout(TIMEOUT_DURATION, tx.send(Ok(request.clone())))
            .await
            .unwrap()
            .unwrap();

        timeout(TIMEOUT_DURATION, tx.send(Ok(listener_request.clone())))
            .await
            .unwrap()
            .unwrap();

        let message = timeout(TIMEOUT_DURATION, stream.next())
            .await
            .unwrap()
            .unwrap()
            .unwrap();
        listener_response.control_plane = message.control_plane.clone();
        listener_response.version_info = message.version_info.clone();
        listener_response.nonce = message.nonce.clone();
        listener_request.response_nonce = message.nonce.clone();

        assert_eq!(listener_response, message);

        timeout(TIMEOUT_DURATION, tx.send(Ok(listener_request.clone())))
            .await
            .unwrap()
            .unwrap();
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::convert::TryFrom;
use std::fs::{remove_file, File, OpenOptions};
use std::io::{Read, Write};

use rustc::ty::layout::Size;

use crate::stacked_borrows::Tag;
use crate::*;

#[derive(Debug)]
pub struct FileHandle {
    file: File,
}

pub struct FileHandler {
    handles: HashMap<i32, FileHandle>,
    low: i32,
}

impl Default for FileHandler {
    fn default() -> Self {
        FileHandler {
            handles: Default::default(),
            // 0, 1 and 2 are reserved for stdin, stdout and stderr.
            low: 3,
        }
    }
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn open(
        &mut self,
        path_op: OpTy<'tcx, Tag>,
        flag_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("open")?;

        let flag = this.read_scalar(flag_op)?.to_i32()?;

        let mut options = OpenOptions::new();

        let o_rdonly = this.eval_libc_i32("O_RDONLY")?;
        let o_wronly = this.eval_libc_i32("O_WRONLY")?;
        let o_rdwr = this.eval_libc_i32("O_RDWR")?;
        // The first two bits of the flag correspond to the access mode in linux, macOS and
        // windows. We need to check that in fact the access mode flags for the current platform
        // only use these two bits, otherwise we are in an unsupported platform and should error.
        if (o_rdonly | o_wronly | o_rdwr) & !0b11 != 0 {
            throw_unsup_format!("Access mode flags on this platform are unsupported");
        }
        // Now we check the access mode
        let access_mode = flag & 0b11;

        if access_mode == o_rdonly {
            options.read(true);
        } else if access_mode == o_wronly {
            options.write(true);
        } else if access_mode == o_rdwr {
            options.read(true).write(true);
        } else {
            throw_unsup_format!("Unsupported access mode {:#x}", access_mode);
        }
        // We need to check that there aren't unsupported options in `flag`. For this we try to
        // reproduce the content of `flag` in the `mirror` variable using only the supported
        // options.
        let mut mirror = access_mode;

        let o_append = this.eval_libc_i32("O_APPEND")?;
        if flag & o_append != 0 {
            options.append(true);
            mirror |= o_append;
        }
        let o_trunc = this.eval_libc_i32("O_TRUNC")?;
        if flag & o_trunc != 0 {
            options.truncate(true);
            mirror |= o_trunc;
        }
        let o_creat = this.eval_libc_i32("O_CREAT")?;
        if flag & o_creat != 0 {
            options.create(true);
            mirror |= o_creat;
        }
        let o_cloexec = this.eval_libc_i32("O_CLOEXEC")?;
        if flag & o_cloexec != 0 {
            // We do not need to do anything for this flag because `std` already sets it.
            // (Technically we do not support *not* setting this flag, but we ignore that.)
            mirror |= o_cloexec;
        }
        // If `flag` is not equal to `mirror`, there is an unsupported option enabled in `flag`,
        // then we throw an error.
        if flag != mirror {
            throw_unsup_format!("unsupported flags {:#x}", flag & !mirror);
        }

        let path = this.read_os_string_from_c_string(this.read_scalar(path_op)?.not_undef()?)?;

        let fd = options.open(path).map(|file| {
            let mut fh = &mut this.machine.file_handler;
            fh.low += 1;
            fh.handles.insert(fh.low, FileHandle { file }).unwrap_none();
            fh.low
        });

        this.try_unwrap_io_result(fd)
    }

    fn fcntl(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        cmd_op: OpTy<'tcx, Tag>,
        _arg1_op: Option<OpTy<'tcx, Tag>>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("fcntl")?;

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let cmd = this.read_scalar(cmd_op)?.to_i32()?;
        // We only support getting the flags for a descriptor.
        if cmd == this.eval_libc_i32("F_GETFD")? {
            // Currently this is the only flag that `F_GETFD` returns. It is OK to just return the
            // `FD_CLOEXEC` value without checking if the flag is set for the file because `std`
            // always sets this flag when opening a file. However we still need to check that the
            // file itself is open.
            if this.machine.file_handler.handles.contains_key(&fd) {
                Ok(this.eval_libc_i32("FD_CLOEXEC")?)
            } else {
                this.handle_not_found()
            }
        } else {
            throw_unsup_format!("The {:#x} command is not supported for `fcntl`)", cmd);
        }
    }

    fn close(&mut self, fd_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("close")?;

        let fd = this.read_scalar(fd_op)?.to_i32()?;

        if let Some(handle) = this.machine.file_handler.handles.remove(&fd) {
            // `File::sync_all` does the checks that are done when closing a file. We do this to
            // to handle possible errors correctly.
            let result = this.try_unwrap_io_result(handle.file.sync_all().map(|_| 0i32));
            // Now we actually close the file.
            drop(handle);
            // And return the result.
            result
        } else {
            this.handle_not_found()
        }
    }

    fn read(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        this.check_no_isolation("read")?;

        let ptr_size = this.pointer_size().bits();

        // We cap the number of read bytes to the largest value that we are able to fit in both the
        // host's and target's `isize`.
        let count = this
            .read_scalar(count_op)?
            .to_machine_usize(&*this.tcx)?
            .min((1 << (ptr_size - 1)) - 1) // max value of target `isize`
            .min(isize::max_value() as u64);

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.read_scalar(buf_op)?.not_undef()?;

        if let Some(handle) = this.machine.file_handler.handles.get_mut(&fd) {
            // This can never fail because `count` was capped to be smaller than
            // `isize::max_value()`.
            let count = isize::try_from(count).unwrap();
            // We want to read at most `count` bytes. We are sure that `count` is not negative
            // because it was a target's `usize`. Also we are sure that its smaller than
            // `usize::max_value()` because it is a host's `isize`.
            let mut bytes = vec![0; count as usize];
            let result = handle
                .file
                .read(&mut bytes)
                // `File::read` never returns a value larger than `count`, so this cannot fail.
                .map(|c| i64::try_from(c).unwrap());

            match result {
                Ok(read_bytes) => {
                    // If reading to `bytes` did not fail, we write those bytes to the buffer.
                    this.memory.write_bytes(buf, bytes)?;
                    Ok(read_bytes)
                }
                Err(e) => {
                    this.set_last_error_from_io_error(e)?;
                    Ok(-1)
                }
            }
        } else {
            this.handle_not_found()
        }
    }

    fn write(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        this.check_no_isolation("write")?;

        let ptr_size = this.pointer_size().bits();

        // We cap the number of read bytes to the largest value that we are able to fit in both the
        // host's and target's `isize`.
        let count = this
            .read_scalar(count_op)?
            .to_machine_usize(&*this.tcx)?
            .min((1 << (ptr_size - 1)) - 1) // max value of target `isize`
            .min(isize::max_value() as u64);

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.read_scalar(buf_op)?.not_undef()?;

        if let Some(handle) = this.machine.file_handler.handles.get_mut(&fd) {
            let bytes = this.memory.read_bytes(buf, Size::from_bytes(count))?;
            let result = handle.file.write(&bytes).map(|c| i64::try_from(c).unwrap());
            this.try_unwrap_io_result(result)
        } else {
            this.handle_not_found()
        }
    }

    fn unlink(&mut self, path_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("unlink")?;

        let path = this.read_os_string_from_c_string(this.read_scalar(path_op)?.not_undef()?)?;

        let result = remove_file(path).map(|_| 0);

        this.try_unwrap_io_result(result)
    }

    /// Function used when a handle is not found inside `FileHandler`. It returns `Ok(-1)`and sets
    /// the last OS error to `libc::EBADF` (invalid file descriptor). This function uses
    /// `T: From<i32>` instead of `i32` directly because some fs functions return different integer
    /// types (like `read`, that returns an `i64`).
    fn handle_not_found<T: From<i32>>(&mut self) -> InterpResult<'tcx, T> {
        let this = self.eval_context_mut();
        let ebadf = this.eval_libc("EBADF")?;
        this.set_last_error(ebadf)?;
        Ok((-1).into())
    }
}
<commit_msg>use new isize_max method in FS accesses; also check full buffers for validity<commit_after>use std::collections::HashMap;
use std::convert::TryFrom;
use std::fs::{remove_file, File, OpenOptions};
use std::io::{Read, Write};

use rustc::ty::layout::{Size, Align};

use crate::stacked_borrows::Tag;
use crate::*;

#[derive(Debug)]
pub struct FileHandle {
    file: File,
}

pub struct FileHandler {
    handles: HashMap<i32, FileHandle>,
    low: i32,
}

impl Default for FileHandler {
    fn default() -> Self {
        FileHandler {
            handles: Default::default(),
            // 0, 1 and 2 are reserved for stdin, stdout and stderr.
            low: 3,
        }
    }
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn open(
        &mut self,
        path_op: OpTy<'tcx, Tag>,
        flag_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("open")?;

        let flag = this.read_scalar(flag_op)?.to_i32()?;

        let mut options = OpenOptions::new();

        let o_rdonly = this.eval_libc_i32("O_RDONLY")?;
        let o_wronly = this.eval_libc_i32("O_WRONLY")?;
        let o_rdwr = this.eval_libc_i32("O_RDWR")?;
        // The first two bits of the flag correspond to the access mode in linux, macOS and
        // windows. We need to check that in fact the access mode flags for the current platform
        // only use these two bits, otherwise we are in an unsupported platform and should error.
        if (o_rdonly | o_wronly | o_rdwr) & !0b11 != 0 {
            throw_unsup_format!("Access mode flags on this platform are unsupported");
        }
        // Now we check the access mode
        let access_mode = flag & 0b11;

        if access_mode == o_rdonly {
            options.read(true);
        } else if access_mode == o_wronly {
            options.write(true);
        } else if access_mode == o_rdwr {
            options.read(true).write(true);
        } else {
            throw_unsup_format!("Unsupported access mode {:#x}", access_mode);
        }
        // We need to check that there aren't unsupported options in `flag`. For this we try to
        // reproduce the content of `flag` in the `mirror` variable using only the supported
        // options.
        let mut mirror = access_mode;

        let o_append = this.eval_libc_i32("O_APPEND")?;
        if flag & o_append != 0 {
            options.append(true);
            mirror |= o_append;
        }
        let o_trunc = this.eval_libc_i32("O_TRUNC")?;
        if flag & o_trunc != 0 {
            options.truncate(true);
            mirror |= o_trunc;
        }
        let o_creat = this.eval_libc_i32("O_CREAT")?;
        if flag & o_creat != 0 {
            options.create(true);
            mirror |= o_creat;
        }
        let o_cloexec = this.eval_libc_i32("O_CLOEXEC")?;
        if flag & o_cloexec != 0 {
            // We do not need to do anything for this flag because `std` already sets it.
            // (Technically we do not support *not* setting this flag, but we ignore that.)
            mirror |= o_cloexec;
        }
        // If `flag` is not equal to `mirror`, there is an unsupported option enabled in `flag`,
        // then we throw an error.
        if flag != mirror {
            throw_unsup_format!("unsupported flags {:#x}", flag & !mirror);
        }

        let path = this.read_os_string_from_c_string(this.read_scalar(path_op)?.not_undef()?)?;

        let fd = options.open(path).map(|file| {
            let mut fh = &mut this.machine.file_handler;
            fh.low += 1;
            fh.handles.insert(fh.low, FileHandle { file }).unwrap_none();
            fh.low
        });

        this.try_unwrap_io_result(fd)
    }

    fn fcntl(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        cmd_op: OpTy<'tcx, Tag>,
        _arg1_op: Option<OpTy<'tcx, Tag>>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("fcntl")?;

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let cmd = this.read_scalar(cmd_op)?.to_i32()?;
        // We only support getting the flags for a descriptor.
        if cmd == this.eval_libc_i32("F_GETFD")? {
            // Currently this is the only flag that `F_GETFD` returns. It is OK to just return the
            // `FD_CLOEXEC` value without checking if the flag is set for the file because `std`
            // always sets this flag when opening a file. However we still need to check that the
            // file itself is open.
            if this.machine.file_handler.handles.contains_key(&fd) {
                Ok(this.eval_libc_i32("FD_CLOEXEC")?)
            } else {
                this.handle_not_found()
            }
        } else {
            throw_unsup_format!("The {:#x} command is not supported for `fcntl`)", cmd);
        }
    }

    fn close(&mut self, fd_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("close")?;

        let fd = this.read_scalar(fd_op)?.to_i32()?;

        if let Some(handle) = this.machine.file_handler.handles.remove(&fd) {
            // `File::sync_all` does the checks that are done when closing a file. We do this to
            // to handle possible errors correctly.
            let result = this.try_unwrap_io_result(handle.file.sync_all().map(|_| 0i32));
            // Now we actually close the file.
            drop(handle);
            // And return the result.
            result
        } else {
            this.handle_not_found()
        }
    }

    fn read(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        this.check_no_isolation("read")?;

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.read_scalar(buf_op)?.not_undef()?;
        let count = this
            .read_scalar(count_op)?
            .to_machine_usize(&*this.tcx)?;

        // Check that the *entire* buffer is actually valid memory.
        this.memory.check_ptr_access(buf, Size::from_bytes(count), Align::from_bytes(1).unwrap())?;

        // We cap the number of read bytes to the largest value that we are able to fit in both the
        // host's and target's `isize`. This saves us from having to handle overflows later.
        let count = count
            .min(this.isize_max() as u64)
            .min(isize::max_value() as u64);

        if let Some(handle) = this.machine.file_handler.handles.get_mut(&fd) {
            // This can never fail because `count` was capped to be smaller than
            // `isize::max_value()`.
            let count = isize::try_from(count).unwrap();
            // We want to read at most `count` bytes. We are sure that `count` is not negative
            // because it was a target's `usize`. Also we are sure that its smaller than
            // `usize::max_value()` because it is a host's `isize`.
            let mut bytes = vec![0; count as usize];
            let result = handle
                .file
                .read(&mut bytes)
                // `File::read` never returns a value larger than `count`, so this cannot fail.
                .map(|c| i64::try_from(c).unwrap());

            match result {
                Ok(read_bytes) => {
                    // If reading to `bytes` did not fail, we write those bytes to the buffer.
                    this.memory.write_bytes(buf, bytes)?;
                    Ok(read_bytes)
                }
                Err(e) => {
                    this.set_last_error_from_io_error(e)?;
                    Ok(-1)
                }
            }
        } else {
            this.handle_not_found()
        }
    }

    fn write(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        this.check_no_isolation("write")?;

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.read_scalar(buf_op)?.not_undef()?;
        let count = this
            .read_scalar(count_op)?
            .to_machine_usize(&*this.tcx)?;

        // Check that the *entire* buffer is actually valid memory.
        this.memory.check_ptr_access(buf, Size::from_bytes(count), Align::from_bytes(1).unwrap())?;

        // We cap the number of written bytes to the largest value that we are able to fit in both the
        // host's and target's `isize`. This saves us from having to handle overflows later.
        let count = count
            .min(this.isize_max() as u64)
            .min(isize::max_value() as u64);

        if let Some(handle) = this.machine.file_handler.handles.get_mut(&fd) {
            let bytes = this.memory.read_bytes(buf, Size::from_bytes(count))?;
            let result = handle.file.write(&bytes).map(|c| i64::try_from(c).unwrap());
            this.try_unwrap_io_result(result)
        } else {
            this.handle_not_found()
        }
    }

    fn unlink(&mut self, path_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        this.check_no_isolation("unlink")?;

        let path = this.read_os_string_from_c_string(this.read_scalar(path_op)?.not_undef()?)?;

        let result = remove_file(path).map(|_| 0);

        this.try_unwrap_io_result(result)
    }

    /// Function used when a handle is not found inside `FileHandler`. It returns `Ok(-1)`and sets
    /// the last OS error to `libc::EBADF` (invalid file descriptor). This function uses
    /// `T: From<i32>` instead of `i32` directly because some fs functions return different integer
    /// types (like `read`, that returns an `i64`).
    fn handle_not_found<T: From<i32>>(&mut self) -> InterpResult<'tcx, T> {
        let this = self.eval_context_mut();
        let ebadf = this.eval_libc("EBADF")?;
        this.set_last_error(ebadf)?;
        Ok((-1).into())
    }
}
<|endoftext|>


<commit_before>//! # Sink module.
//!
//! The Sink module send metrics to Warp10.
use std::collections::VecDeque;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;

use hyper;
use hyper_tls::HttpsConnector;
use hyper_timeout::TimeoutConnector;

use tokio_core::reactor::Core;

use futures::future::Shared;
use futures::Future;
use futures::sync::oneshot;
use futures::sync::mpsc::{channel, Sender};
use futures::task::Task;

use slog_scope;

use config;

mod fs;
mod send;

const REACTOR_CLIENT: usize = 30;

pub struct Sink<'a> {
    name: &'a String,
    todo: Arc<Mutex<VecDeque<PathBuf>>>,
    handles: Vec<thread::JoinHandle<()>>,
    dir: &'a String,
    watch_period: u64,
    timeout: u64,
    keep_alive: bool,
    token: &'a String,
    token_header: &'a String,
    url: &'a hyper::Uri,
    batch_count: u64,
    batch_size: u64,
    max_size: u64,
    ttl: u64,
    parallel: u64,
    sigint: (oneshot::Sender<()>, Shared<oneshot::Receiver<()>>),
}

impl<'a> Sink<'a> {
    pub fn new(sink: &'a config::Sink, parameters: &'a config::Parameters) -> Sink<'a> {
        let (shutdown_tx, shutdown_rx) = oneshot::channel();

        Sink {
            name: &sink.name,
            todo: Arc::new(Mutex::new(VecDeque::new())),
            handles: Vec::new(),
            sigint: (shutdown_tx, shutdown_rx.shared()),
            dir: &parameters.sink_dir,
            watch_period: parameters.scan_period,
            timeout: parameters.timeout,
            keep_alive: sink.keep_alive,
            token: &sink.token,
            token_header: &sink.token_header,
            url: &sink.url,
            batch_count: parameters.batch_count,
            batch_size: parameters.batch_size,
            max_size: sink.size,
            ttl: sink.ttl,
            parallel: sink.parallel,
        }
    }

    pub fn start(&mut self) {
        debug!("start sink: {}", self.name);
        let (notify_tx, notify_rx) = channel(self.parallel as usize);

        // spawn fs thread
        let (name, dir, todo) = (self.name.clone(), self.dir.clone(), self.todo.clone());
        let (period, sigint) = (self.watch_period, self.sigint.1.clone());
        let (max_size, ttl) = (self.max_size, self.ttl);

        self.handles.push(thread::spawn(move || {
            slog_scope::scope(
                &slog_scope::logger().new(o!("sink" => name.clone())),
                || fs::fs_thread(&name, &dir, period, todo, max_size, ttl, sigint, notify_rx),
            );
        }));

        // spawn sender threads
        let reactor_count = (self.parallel as f64 / REACTOR_CLIENT as f64).ceil() as u64;
        let client_count = (self.parallel as f64 / reactor_count as f64).ceil() as u64;
        for _ in 0..reactor_count {
            let (name, sigint) = (self.name.clone(), self.sigint.1.clone());
            let todo = self.todo.clone();
            let url = self.url.clone();
            let (timeout, keep_alive) = (self.timeout, self.keep_alive);
            let (token, token_header) = (self.token.clone(), self.token_header.clone());
            let (batch_count, batch_size) = (self.batch_count, self.batch_size);
            let notify_tx: Sender<Task> = notify_tx.clone();

            self.handles.push(thread::spawn(move || {
                slog_scope::scope(
                    &slog_scope::logger().new(o!("sink" => name.clone())),
                    || {
                        let mut core = Core::new().expect("Fail to start tokio reactor");
                        let handle = core.handle();
                        let connector =
                            HttpsConnector::new(4, &handle).expect("Fail to start https connector");

                        // Handle connection timeouts
                        let mut tm = TimeoutConnector::new(connector, &handle);
                        tm.set_connect_timeout(Some(Duration::from_secs(timeout)));
                        tm.set_read_timeout(Some(Duration::from_secs(timeout)));
                        tm.set_write_timeout(Some(Duration::from_secs(timeout)));

                        let client = Arc::new(
                            hyper::Client::configure()
                                .body::<send::PayloadBody>()
                                .keep_alive(keep_alive)
                                .connector(tm)
                                .build(&handle),
                        );

                        // spawn send threads
                        for _ in 0..client_count {
                            let work = send::send_thread(
                                client.clone(),
                                token.clone(),
                                token_header.clone(),
                                url.clone(),
                                todo.clone(),
                                batch_count,
                                batch_size,
                                notify_tx.clone(),
                            );

                            core.handle().spawn(work);
                        }

                        core.run(sigint).expect("Sigint could not fail");
                    },
                )
            }));
        }
    }

    pub fn stop(self) {
        self.sigint.0.send(()).unwrap();
        for handle in self.handles {
            handle.join().unwrap();
        }
    }
}
<commit_msg>fix(sink): remove read/write timeout<commit_after>//! # Sink module.
//!
//! The Sink module send metrics to Warp10.
use std::collections::VecDeque;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;

use hyper;
use hyper_tls::HttpsConnector;
use hyper_timeout::TimeoutConnector;

use tokio_core::reactor::Core;

use futures::future::Shared;
use futures::Future;
use futures::sync::oneshot;
use futures::sync::mpsc::{channel, Sender};
use futures::task::Task;

use slog_scope;

use config;

mod fs;
mod send;

const REACTOR_CLIENT: usize = 30;

pub struct Sink<'a> {
    name: &'a String,
    todo: Arc<Mutex<VecDeque<PathBuf>>>,
    handles: Vec<thread::JoinHandle<()>>,
    dir: &'a String,
    watch_period: u64,
    timeout: u64,
    keep_alive: bool,
    token: &'a String,
    token_header: &'a String,
    url: &'a hyper::Uri,
    batch_count: u64,
    batch_size: u64,
    max_size: u64,
    ttl: u64,
    parallel: u64,
    sigint: (oneshot::Sender<()>, Shared<oneshot::Receiver<()>>),
}

impl<'a> Sink<'a> {
    pub fn new(sink: &'a config::Sink, parameters: &'a config::Parameters) -> Sink<'a> {
        let (shutdown_tx, shutdown_rx) = oneshot::channel();

        Sink {
            name: &sink.name,
            todo: Arc::new(Mutex::new(VecDeque::new())),
            handles: Vec::new(),
            sigint: (shutdown_tx, shutdown_rx.shared()),
            dir: &parameters.sink_dir,
            watch_period: parameters.scan_period,
            timeout: parameters.timeout,
            keep_alive: sink.keep_alive,
            token: &sink.token,
            token_header: &sink.token_header,
            url: &sink.url,
            batch_count: parameters.batch_count,
            batch_size: parameters.batch_size,
            max_size: sink.size,
            ttl: sink.ttl,
            parallel: sink.parallel,
        }
    }

    pub fn start(&mut self) {
        debug!("start sink: {}", self.name);
        let (notify_tx, notify_rx) = channel(self.parallel as usize);

        // spawn fs thread
        let (name, dir, todo) = (self.name.clone(), self.dir.clone(), self.todo.clone());
        let (period, sigint) = (self.watch_period, self.sigint.1.clone());
        let (max_size, ttl) = (self.max_size, self.ttl);

        self.handles.push(thread::spawn(move || {
            slog_scope::scope(
                &slog_scope::logger().new(o!("sink" => name.clone())),
                || fs::fs_thread(&name, &dir, period, todo, max_size, ttl, sigint, notify_rx),
            );
        }));

        // spawn sender threads
        let reactor_count = (self.parallel as f64 / REACTOR_CLIENT as f64).ceil() as u64;
        let client_count = (self.parallel as f64 / reactor_count as f64).ceil() as u64;
        for _ in 0..reactor_count {
            let (name, sigint) = (self.name.clone(), self.sigint.1.clone());
            let todo = self.todo.clone();
            let url = self.url.clone();
            let (timeout, keep_alive) = (self.timeout, self.keep_alive);
            let (token, token_header) = (self.token.clone(), self.token_header.clone());
            let (batch_count, batch_size) = (self.batch_count, self.batch_size);
            let notify_tx: Sender<Task> = notify_tx.clone();

            self.handles.push(thread::spawn(move || {
                slog_scope::scope(
                    &slog_scope::logger().new(o!("sink" => name.clone())),
                    || {
                        let mut core = Core::new().expect("Fail to start tokio reactor");
                        let handle = core.handle();
                        let connector =
                            HttpsConnector::new(4, &handle).expect("Fail to start https connector");

                        // Handle connection timeouts
                        let mut tm = TimeoutConnector::new(connector, &handle);
                        tm.set_connect_timeout(Some(Duration::from_secs(timeout)));

                        let client = Arc::new(
                            hyper::Client::configure()
                                .body::<send::PayloadBody>()
                                .keep_alive(keep_alive)
                                .connector(tm)
                                .build(&handle),
                        );

                        // spawn send threads
                        for _ in 0..client_count {
                            let work = send::send_thread(
                                client.clone(),
                                token.clone(),
                                token_header.clone(),
                                url.clone(),
                                todo.clone(),
                                batch_count,
                                batch_size,
                                notify_tx.clone(),
                            );

                            core.handle().spawn(work);
                        }

                        core.run(sigint).expect("Sigint could not fail");
                    },
                )
            }));
        }
    }

    pub fn stop(self) {
        self.sigint.0.send(()).unwrap();
        for handle in self.handles {
            handle.join().unwrap();
        }
    }
}
<|endoftext|>


<commit_before>use std::ops;
use std::mem;
use std::ptr;
use std::ptr::Unique;
use std::marker::Unsize;
use std::fmt;
use std::hash;
use std::hash::Hash;
use std::cmp::Ordering;
use std::ops::CoerceUnsized;
use nodrop_union::NoDrop;

use super::space::S2;

impl<T: ?Sized + Unsize<U>, U: ?Sized> CoerceUnsized<StackBox<U>> for StackBox<T> {}

/// On-stack allocation for dynamically-sized type.
///
/// # Examples
///
/// ```
/// use smallbox::StackBox;
///
/// let val: StackBox<PartialEq<usize>> = StackBox::new(5usize).unwrap();
///
/// assert!(*val == 5)
/// ```
pub struct StackBox<T: ?Sized, Space = S2> {
    ptr: Unique<T>,
    space: NoDrop<Space>,
}

impl<T: ?Sized, Space> StackBox<T, Space> {
    /// Try to alloc on stack, and return Err<T>
    /// when val is larger than capacity
    ///
    /// # Examples
    ///
    /// ```
    /// use smallbox::StackBox;
    ///
    /// assert!(StackBox::<[usize]>::new([0usize; 1]).is_ok());
    /// assert!(StackBox::<[usize]>::new([0usize; 8]).is_err());
    /// ```
    pub fn new<U>(mut val: U) -> Result<StackBox<T, Space>, U>
        where U: Unsize<T> + Sized
    {
        if mem::size_of::<U>() > mem::size_of::<Space>() {
            Err(val)
        } else {
            unsafe {
                let ptr: Unique<T> = Unique::new(&mut val).unwrap();
                let mut space = NoDrop::new(mem::uninitialized::<Space>());
                ptr::copy_nonoverlapping(&val, &mut space as *mut _ as *mut U, 1);
                mem::forget(val);
                Ok(StackBox { ptr, space })
            }
        }
    }

    /// Try to transforms to a `StackBox<T>` with bigger capacity,
    /// and return `Err` when target capacity is smaller.
    ///
    /// # Example
    ///
    /// ```
    /// use smallbox::StackBox;
    /// use smallbox::space::*;
    ///
    /// let s = StackBox::<[usize], S8>::new([0usize; 4]).unwrap();
    /// assert!(s.resize::<S16>().is_ok());
    ///
    /// let s = StackBox::<[usize], S8>::new([0usize; 4]).unwrap();
    /// assert!(s.resize::<S4>().is_err());
    /// ```
    pub fn resize<ToSpace>(self) -> Result<StackBox<T, ToSpace>, Self> {
        if mem::size_of::<Space>() > mem::size_of::<ToSpace>() {
            Err(self)
        } else {
            unsafe {
                let ptr = self.ptr;
                let mut space = NoDrop::new(mem::uninitialized::<ToSpace>());
                ptr::copy_nonoverlapping(&self.space,
                                         &mut space as *mut _ as *mut NoDrop<Space>,
                                         1);
                mem::forget(self);
                Ok(StackBox { ptr, space })
            }
        }
    }

    unsafe fn as_ptr(&self) -> *const T {
        let mut ptr: *const T = self.ptr.as_ptr();
        *(&mut ptr as *mut _ as *mut usize) = &self.space as *const _ as usize;
        ptr
    }
}

impl<T: ?Sized, Space> ops::Deref for StackBox<T, Space> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { &*self.as_ptr() }
    }
}

impl<T: ?Sized, Space> ops::DerefMut for StackBox<T, Space> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *(self.as_ptr() as *mut T) }
    }
}

impl<T: ?Sized, Space> ops::Drop for StackBox<T, Space> {
    fn drop(&mut self) {
        unsafe { ptr::drop_in_place(&mut **self) }
    }
}

impl<T: fmt::Display + ?Sized, Space> fmt::Display for StackBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&**self, f)
    }
}

impl<T: fmt::Debug + ?Sized, Space> fmt::Debug for StackBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&**self, f)
    }
}

impl<T: ?Sized, Space> fmt::Pointer for StackBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // It's not possible to extract the inner Uniq directly from the Box,
        // instead we cast it to a *const which aliases the Unique
        let ptr: *const T = &**self;
        fmt::Pointer::fmt(&ptr, f)
    }
}

impl<T: ?Sized + PartialEq, Space> PartialEq for StackBox<T, Space> {
    #[inline]
    fn eq(&self, other: &StackBox<T, Space>) -> bool {
        PartialEq::eq(&**self, &**other)
    }
    #[inline]
    fn ne(&self, other: &StackBox<T, Space>) -> bool {
        PartialEq::ne(&**self, &**other)
    }
}

impl<T: ?Sized + PartialOrd, Space> PartialOrd for StackBox<T, Space> {
    #[inline]
    fn partial_cmp(&self, other: &StackBox<T, Space>) -> Option<Ordering> {
        PartialOrd::partial_cmp(&**self, &**other)
    }
    #[inline]
    fn lt(&self, other: &StackBox<T, Space>) -> bool {
        PartialOrd::lt(&**self, &**other)
    }
    #[inline]
    fn le(&self, other: &StackBox<T, Space>) -> bool {
        PartialOrd::le(&**self, &**other)
    }
    #[inline]
    fn ge(&self, other: &StackBox<T, Space>) -> bool {
        PartialOrd::ge(&**self, &**other)
    }
    #[inline]
    fn gt(&self, other: &StackBox<T, Space>) -> bool {
        PartialOrd::gt(&**self, &**other)
    }
}

impl<T: ?Sized + Ord, Space> Ord for StackBox<T, Space> {
    #[inline]
    fn cmp(&self, other: &StackBox<T, Space>) -> Ordering {
        Ord::cmp(&**self, &**other)
    }
}

impl<T: ?Sized + Eq, Space> Eq for StackBox<T, Space> {}

impl<T: ?Sized + Hash, Space> Hash for StackBox<T, Space> {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        (**self).hash(state);
    }
}<commit_msg>fix Unique breaking change<commit_after>use std::ops;
use std::mem;
use std::ptr;
use std::ptr::Unique;
use std::marker::Unsize;
use std::marker::PhantomData;
use std::fmt;
use std::hash;
use std::hash::Hash;
use std::cmp::Ordering;
use std::ops::CoerceUnsized;
use nodrop_union::NoDrop;

use super::space::S2;

impl<T: ?Sized + Unsize<U>, U: ?Sized> CoerceUnsized<StackBox<U>> for StackBox<T> {}

/// On-stack allocation for dynamically-sized type.
///
/// # Examples
///
/// ```
/// use smallbox::StackBox;
///
/// let val: StackBox<PartialEq<usize>> = StackBox::new(5usize).unwrap();
///
/// assert!(*val == 5)
/// ```
pub struct StackBox<T: ?Sized, Space = S2> {
    ptr: *const T,
    space: NoDrop<Space>,
    _phantom: PhantomData<T>,
}

impl<T: ?Sized, Space> StackBox<T, Space> {
    /// Try to alloc on stack, and return Err<T>
    /// when val is larger than capacity
    ///
    /// # Examples
    ///
    /// ```
    /// use smallbox::StackBox;
    ///
    /// assert!(StackBox::<[usize]>::new([0usize; 1]).is_ok());
    /// assert!(StackBox::<[usize]>::new([0usize; 8]).is_err());
    /// ```
    pub fn new<U>(mut val: U) -> Result<StackBox<T, Space>, U>
        where U: Unsize<T> + Sized
    {
        if mem::size_of::<U>() > mem::size_of::<Space>() {
            Err(val)
        } else {
            unsafe {
                let ptr = &val as *const _;
                let mut space = NoDrop::new(mem::uninitialized::<Space>());
                ptr::copy_nonoverlapping(&val, &mut space as *mut _ as *mut U, 1);
                mem::forget(val);
                Ok(StackBox {
                       ptr,
                       space,
                       _phantom: PhantomData,
                   })
            }
        }
    }

    /// Try to transforms to a `StackBox<T>` with bigger capacity,
    /// and return `Err` when target capacity is smaller.
    ///
    /// # Example
    ///
    /// ```
    /// use smallbox::StackBox;
    /// use smallbox::space::*;
    ///
    /// let s = StackBox::<[usize], S8>::new([0usize; 4]).unwrap();
    /// assert!(s.resize::<S16>().is_ok());
    ///
    /// let s = StackBox::<[usize], S8>::new([0usize; 4]).unwrap();
    /// assert!(s.resize::<S4>().is_err());
    /// ```
    pub fn resize<ToSpace>(self) -> Result<StackBox<T, ToSpace>, Self> {
        if mem::size_of::<Space>() > mem::size_of::<ToSpace>() {
            Err(self)
        } else {
            unsafe {
                let ptr = self.ptr;
                let mut space = NoDrop::new(mem::uninitialized::<ToSpace>());
                ptr::copy_nonoverlapping(&self.space,
                                         &mut space as *mut _ as *mut NoDrop<Space>,
                                         1);
                mem::forget(self);
                Ok(StackBox {
                       ptr,
                       space,
                       _phantom: PhantomData,
                   })
            }
        }
    }

    unsafe fn as_ptr(&self) -> *const T {
        let mut ptr: *const T = self.ptr;
        *(&mut ptr as *mut _ as *mut usize) = &self.space as *const _ as usize;
        ptr
    }
}

impl<T: ?Sized, Space> ops::Deref for StackBox<T, Space> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { &*self.as_ptr() }
    }
}

impl<T: ?Sized, Space> ops::DerefMut for StackBox<T, Space> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *(self.as_ptr() as *mut T) }
    }
}

impl<T: ?Sized, Space> ops::Drop for StackBox<T, Space> {
    fn drop(&mut self) {
        unsafe { ptr::drop_in_place(&mut **self) }
    }
}

impl<T: fmt::Display + ?Sized, Space> fmt::Display for StackBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&**self, f)
    }
}

impl<T: fmt::Debug + ?Sized, Space> fmt::Debug for StackBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&**self, f)
    }
}

impl<T: ?Sized, Space> fmt::Pointer for StackBox<T, Space> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // It's not possible to extract the inner Uniq directly from the Box,
        // instead we cast it to a *const which aliases the Unique
        let ptr: *const T = &**self;
        fmt::Pointer::fmt(&ptr, f)
    }
}

impl<T: ?Sized + PartialEq, Space> PartialEq for StackBox<T, Space> {
    #[inline]
    fn eq(&self, other: &StackBox<T, Space>) -> bool {
        PartialEq::eq(&**self, &**other)
    }
    #[inline]
    fn ne(&self, other: &StackBox<T, Space>) -> bool {
        PartialEq::ne(&**self, &**other)
    }
}

impl<T: ?Sized + PartialOrd, Space> PartialOrd for StackBox<T, Space> {
    #[inline]
    fn partial_cmp(&self, other: &StackBox<T, Space>) -> Option<Ordering> {
        PartialOrd::partial_cmp(&**self, &**other)
    }
    #[inline]
    fn lt(&self, other: &StackBox<T, Space>) -> bool {
        PartialOrd::lt(&**self, &**other)
    }
    #[inline]
    fn le(&self, other: &StackBox<T, Space>) -> bool {
        PartialOrd::le(&**self, &**other)
    }
    #[inline]
    fn ge(&self, other: &StackBox<T, Space>) -> bool {
        PartialOrd::ge(&**self, &**other)
    }
    #[inline]
    fn gt(&self, other: &StackBox<T, Space>) -> bool {
        PartialOrd::gt(&**self, &**other)
    }
}

impl<T: ?Sized + Ord, Space> Ord for StackBox<T, Space> {
    #[inline]
    fn cmp(&self, other: &StackBox<T, Space>) -> Ordering {
        Ord::cmp(&**self, &**other)
    }
}

impl<T: ?Sized + Eq, Space> Eq for StackBox<T, Space> {}

impl<T: ?Sized + Hash, Space> Hash for StackBox<T, Space> {
    fn hash<H: hash::Hasher>(&self, state: &mut H) {
        (**self).hash(state);
    }
}<|endoftext|>