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

use std::fmt::{self, Display, Formatter};
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use futures::future::{Future, TryFutureExt};
use futures::sink::SinkExt;
use futures::stream::{Stream, StreamExt, TryStreamExt};
use futures::task::{Context, Poll};
use grpcio::{
    ChannelBuilder, ClientStreamingSink, Environment, RequestStream, RpcStatus, RpcStatusCode,
    WriteFlags,
};
use kvproto::raft_serverpb::{Done, RaftMessage, RaftSnapshotData, SnapshotChunk};
use kvproto::tikvpb::TikvClient;
use protobuf::Message;
use tokio::runtime::{Builder as RuntimeBuilder, Runtime};

use engine_rocks::RocksEngine;
use file_system::{IOType, WithIOType};
use raftstore::router::RaftStoreRouter;
use raftstore::store::{GenericSnapshot, SnapEntry, SnapKey, SnapManager};
use security::SecurityManager;
use tikv_util::worker::Runnable;
use tikv_util::DeferContext;

use super::metrics::*;
use super::{Config, Error, Result};

pub type Callback = Box<dyn FnOnce(Result<()>) + Send>;

const DEFAULT_POOL_SIZE: usize = 4;

/// A task for either receiving Snapshot or sending Snapshot
pub enum Task {
    Recv {
        stream: RequestStream<SnapshotChunk>,
        sink: ClientStreamingSink<Done>,
    },
    Send {
        addr: String,
        msg: RaftMessage,
        cb: Callback,
    },
}

impl Display for Task {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match *self {
            Task::Recv { .. } => write!(f, "Recv"),
            Task::Send {
                ref addr, ref msg, ..
            } => write!(f, "Send Snap[to: {}, snap: {:?}]", addr, msg),
        }
    }
}

struct SnapChunk {
    first: Option<SnapshotChunk>,
    snap: Box<dyn GenericSnapshot>,
    remain_bytes: usize,
}

const SNAP_CHUNK_LEN: usize = 1024 * 1024;

impl Stream for SnapChunk {
    type Item = Result<(SnapshotChunk, WriteFlags)>;

    fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        if let Some(t) = self.first.take() {
            let write_flags = WriteFlags::default().buffer_hint(true);
            return Poll::Ready(Some(Ok((t, write_flags))));
        }

        let mut buf = match self.remain_bytes {
            0 => return Poll::Ready(None),
            n if n > SNAP_CHUNK_LEN => vec![0; SNAP_CHUNK_LEN],
            n => vec![0; n],
        };
        let result = self.snap.read_exact(buf.as_mut_slice());
        match result {
            Ok(_) => {
                self.remain_bytes -= buf.len();
                let mut chunk = SnapshotChunk::default();
                chunk.set_data(buf);
                Poll::Ready(Some(Ok((chunk, WriteFlags::default().buffer_hint(true)))))
            }
            Err(e) => Poll::Ready(Some(Err(box_err!("failed to read snapshot chunk: {}", e)))),
        }
    }
}

struct SendStat {
    key: SnapKey,
    total_size: u64,
    elapsed: Duration,
}

/// Send the snapshot to specified address.
///
/// It will first send the normal raft snapshot message and then send the snapshot file.
fn send_snap(
    env: Arc<Environment>,
    mgr: SnapManager,
    security_mgr: Arc<SecurityManager>,
    cfg: &Config,
    addr: &str,
    msg: RaftMessage,
) -> Result<impl Future<Output = Result<SendStat>>> {
    assert!(msg.get_message().has_snapshot());
    let timer = Instant::now();

    let send_timer = SEND_SNAP_HISTOGRAM.start_coarse_timer();

    let key = {
        let snap = msg.get_message().get_snapshot();
        SnapKey::from_snap(snap)?
    };

    mgr.register(key.clone(), SnapEntry::Sending);
    let deregister = {
        let (mgr, key) = (mgr.clone(), key.clone());
        DeferContext::new(move || mgr.deregister(&key, &SnapEntry::Sending))
    };

    let s = box_try!(mgr.get_snapshot_for_sending(&key));
    if !s.exists() {
        return Err(box_err!("missing snap file: {:?}", s.path()));
    }
    let total_size = s.total_size()?;

    let mut chunks = {
        let mut first_chunk = SnapshotChunk::default();
        first_chunk.set_message(msg);

        SnapChunk {
            first: Some(first_chunk),
            snap: s,
            remain_bytes: total_size as usize,
        }
    };

    let cb = ChannelBuilder::new(env)
        .stream_initial_window_size(cfg.grpc_stream_initial_window_size.0 as i32)
        .keepalive_time(cfg.grpc_keepalive_time.0)
        .keepalive_timeout(cfg.grpc_keepalive_timeout.0)
        .default_compression_algorithm(cfg.grpc_compression_algorithm());

    let channel = security_mgr.connect(cb, addr);
    let client = TikvClient::new(channel);
    let (sink, receiver) = client.snapshot()?;

    let send_task = async move {
        let mut sink = sink.sink_map_err(Error::from);
        sink.send_all(&mut chunks).await?;
        sink.close().await?;
        let recv_result = receiver.map_err(Error::from).await;
        send_timer.observe_duration();
        drop(deregister);
        drop(client);
        match recv_result {
            Ok(_) => {
                fail_point!("snapshot_delete_after_send");
                chunks.snap.delete();
                // TODO: improve it after rustc resolves the bug.
                // Call `info` in the closure directly will cause rustc
                // panic with `Cannot create local mono-item for DefId`.
                Ok(SendStat {
                    key,
                    total_size,
                    elapsed: timer.elapsed(),
                })
            }
            Err(e) => Err(e),
        }
    };
    Ok(send_task)
}

struct RecvSnapContext {
    key: SnapKey,
    file: Option<Box<dyn GenericSnapshot>>,
    raft_msg: RaftMessage,
    _with_io_type: WithIOType,
}

impl RecvSnapContext {
    fn new(head_chunk: Option<SnapshotChunk>, snap_mgr: &SnapManager) -> Result<Self> {
        // head_chunk is None means the stream is empty.
        let mut head = head_chunk.ok_or_else(|| Error::Other("empty gRPC stream".into()))?;
        if !head.has_message() {
            return Err(box_err!("no raft message in the first chunk"));
        }

        let meta = head.take_message();
        let key = match SnapKey::from_snap(meta.get_message().get_snapshot()) {
            Ok(k) => k,
            Err(e) => return Err(box_err!("failed to create snap key: {:?}", e)),
        };

        let data = meta.get_message().get_snapshot().get_data();
        let mut snapshot = RaftSnapshotData::default();
        snapshot.merge_from_bytes(data)?;
        let with_io_type = WithIOType::new(if snapshot.get_meta().get_for_balance() {
            IOType::LoadBalance
        } else {
            IOType::Replication
        });

        let snap = {
            let s = match snap_mgr.get_snapshot_for_receiving(&key, data) {
                Ok(s) => s,
                Err(e) => return Err(box_err!("{} failed to create snapshot file: {:?}", key, e)),
            };

            if s.exists() {
                let p = s.path();
                info!("snapshot file already exists, skip receiving"; "snap_key" => %key, "file" => p);
                None
            } else {
                Some(s)
            }
        };

        Ok(RecvSnapContext {
            key,
            file: snap,
            raft_msg: meta,
            _with_io_type: with_io_type,
        })
    }

    fn finish<R: RaftStoreRouter<RocksEngine>>(self, raft_router: R) -> Result<()> {
        let key = self.key;
        if let Some(mut file) = self.file {
            info!("saving snapshot file"; "snap_key" => %key, "file" => file.path());
            if let Err(e) = file.save() {
                let path = file.path();
                let e = box_err!("{} failed to save snapshot file {}: {:?}", key, path, e);
                return Err(e);
            }
        }
        if let Err(e) = raft_router.send_raft_msg(self.raft_msg) {
            return Err(box_err!("{} failed to send snapshot to raft: {}", key, e));
        }
        Ok(())
    }
}

fn recv_snap<R: RaftStoreRouter<RocksEngine> + 'static>(
    stream: RequestStream<SnapshotChunk>,
    sink: ClientStreamingSink<Done>,
    snap_mgr: SnapManager,
    raft_router: R,
) -> impl Future<Output = Result<()>> {
    let recv_task = async move {
        let mut stream = stream.map_err(Error::from);
        let head = stream.next().await.transpose()?;
        let mut context = RecvSnapContext::new(head, &snap_mgr)?;
        if context.file.is_none() {
            return context.finish(raft_router);
        }
        let context_key = context.key.clone();
        snap_mgr.register(context.key.clone(), SnapEntry::Receiving);

        while let Some(item) = stream.next().await {
            let mut chunk = item?;
            let data = chunk.take_data();
            if data.is_empty() {
                return Err(box_err!("{} receive chunk with empty data", context.key));
            }
            if let Err(e) = context.file.as_mut().unwrap().write_all(&data) {
                let key = &context.key;
                let path = context.file.as_mut().unwrap().path();
                let e = box_err!("{} failed to write snapshot file {}: {}", key, path, e);
                return Err(e);
            }
        }

        let res = context.finish(raft_router);
        snap_mgr.deregister(&context_key, &SnapEntry::Receiving);
        res
    };

    async move {
        match recv_task.await {
            Ok(()) => sink.success(Done::default()).await.map_err(Error::from),
            Err(e) => {
                let status = RpcStatus::new(RpcStatusCode::UNKNOWN, Some(format!("{:?}", e)));
                sink.fail(status).await.map_err(Error::from)
            }
        }
    }
}

pub struct Runner<R: RaftStoreRouter<RocksEngine> + 'static> {
    env: Arc<Environment>,
    snap_mgr: SnapManager,
    pool: Runtime,
    raft_router: R,
    security_mgr: Arc<SecurityManager>,
    cfg: Arc<Config>,
    sending_count: Arc<AtomicUsize>,
    recving_count: Arc<AtomicUsize>,
}

impl<R: RaftStoreRouter<RocksEngine> + 'static> Runner<R> {
    pub fn new(
        env: Arc<Environment>,
        snap_mgr: SnapManager,
        r: R,
        security_mgr: Arc<SecurityManager>,
        cfg: Arc<Config>,
    ) -> Runner<R> {
        Runner {
            env,
            snap_mgr,
            pool: RuntimeBuilder::new()
                .threaded_scheduler()
                .thread_name(thd_name!("snap-sender"))
                .core_threads(DEFAULT_POOL_SIZE)
                .on_thread_start(tikv_alloc::add_thread_memory_accessor)
                .on_thread_stop(tikv_alloc::remove_thread_memory_accessor)
                .build()
                .unwrap(),
            raft_router: r,
            security_mgr,
            cfg,
            sending_count: Arc::new(AtomicUsize::new(0)),
            recving_count: Arc::new(AtomicUsize::new(0)),
        }
    }
}

impl<R: RaftStoreRouter<RocksEngine> + 'static> Runnable for Runner<R> {
    type Task = Task;

    fn run(&mut self, task: Task) {
        match task {
            Task::Recv { stream, sink } => {
                let task_num = self.recving_count.load(Ordering::SeqCst);
                if task_num >= self.cfg.concurrent_recv_snap_limit {
                    warn!("too many recving snapshot tasks, ignore");
                    let status = RpcStatus::new(
                        RpcStatusCode::RESOURCE_EXHAUSTED,
                        Some(format!(
                            "the number of received snapshot tasks {} exceeded the limitation {}",
                            task_num, self.cfg.concurrent_recv_snap_limit
                        )),
                    );
                    self.pool.spawn(sink.fail(status));
                    return;
                }
                SNAP_TASK_COUNTER_STATIC.recv.inc();

                let snap_mgr = self.snap_mgr.clone();
                let raft_router = self.raft_router.clone();
                let recving_count = Arc::clone(&self.recving_count);
                recving_count.fetch_add(1, Ordering::SeqCst);
                let task = async move {
                    let result = recv_snap(stream, sink, snap_mgr, raft_router).await;
                    recving_count.fetch_sub(1, Ordering::SeqCst);
                    if let Err(e) = result {
                        error!("failed to recv snapshot"; "err" => %e);
                    }
                };
                self.pool.spawn(task);
            }
            Task::Send { addr, msg, cb } => {
                fail_point!("send_snapshot");
                if self.sending_count.load(Ordering::SeqCst) >= self.cfg.concurrent_send_snap_limit
                {
                    warn!(
                        "too many sending snapshot tasks, drop Send Snap[to: {}, snap: {:?}]",
                        addr, msg
                    );
                    cb(Err(Error::Other("Too many sending snapshot tasks".into())));
                    return;
                }
                SNAP_TASK_COUNTER_STATIC.send.inc();

                let env = Arc::clone(&self.env);
                let mgr = self.snap_mgr.clone();
                let security_mgr = Arc::clone(&self.security_mgr);
                let sending_count = Arc::clone(&self.sending_count);
                sending_count.fetch_add(1, Ordering::SeqCst);

                let send_task = send_snap(env, mgr, security_mgr, &self.cfg, &addr, msg);
                let task = async move {
                    let res = match send_task {
                        Err(e) => Err(e),
                        Ok(f) => f.await,
                    };
                    match res {
                        Ok(stat) => {
                            info!(
                                "sent snapshot";
                                "region_id" => stat.key.region_id,
                                "snap_key" => %stat.key,
                                "size" => stat.total_size,
                                "duration" => ?stat.elapsed
                            );
                            cb(Ok(()));
                        }
                        Err(e) => {
                            error!("failed to send snap"; "to_addr" => addr, "err" => ?e);
                            cb(Err(e));
                        }
                    };
                    sending_count.fetch_sub(1, Ordering::SeqCst);
                };

                self.pool.spawn(task);
            }
        }
    }
}
<commit_msg>Make the server::snap module generic over engines (#10014)<commit_after>// Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.

use std::fmt::{self, Display, Formatter};
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use futures::future::{Future, TryFutureExt};
use futures::sink::SinkExt;
use futures::stream::{Stream, StreamExt, TryStreamExt};
use futures::task::{Context, Poll};
use grpcio::{
    ChannelBuilder, ClientStreamingSink, Environment, RequestStream, RpcStatus, RpcStatusCode,
    WriteFlags,
};
use kvproto::raft_serverpb::{Done, RaftMessage, RaftSnapshotData, SnapshotChunk};
use kvproto::tikvpb::TikvClient;
use protobuf::Message;
use tokio::runtime::{Builder as RuntimeBuilder, Runtime};

use engine_traits::KvEngine;
use file_system::{IOType, WithIOType};
use raftstore::router::RaftStoreRouter;
use raftstore::store::{GenericSnapshot, SnapEntry, SnapKey, SnapManager};
use security::SecurityManager;
use tikv_util::worker::Runnable;
use tikv_util::DeferContext;

use super::metrics::*;
use super::{Config, Error, Result};

pub type Callback = Box<dyn FnOnce(Result<()>) + Send>;

const DEFAULT_POOL_SIZE: usize = 4;

/// A task for either receiving Snapshot or sending Snapshot
pub enum Task {
    Recv {
        stream: RequestStream<SnapshotChunk>,
        sink: ClientStreamingSink<Done>,
    },
    Send {
        addr: String,
        msg: RaftMessage,
        cb: Callback,
    },
}

impl Display for Task {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match *self {
            Task::Recv { .. } => write!(f, "Recv"),
            Task::Send {
                ref addr, ref msg, ..
            } => write!(f, "Send Snap[to: {}, snap: {:?}]", addr, msg),
        }
    }
}

struct SnapChunk {
    first: Option<SnapshotChunk>,
    snap: Box<dyn GenericSnapshot>,
    remain_bytes: usize,
}

const SNAP_CHUNK_LEN: usize = 1024 * 1024;

impl Stream for SnapChunk {
    type Item = Result<(SnapshotChunk, WriteFlags)>;

    fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        if let Some(t) = self.first.take() {
            let write_flags = WriteFlags::default().buffer_hint(true);
            return Poll::Ready(Some(Ok((t, write_flags))));
        }

        let mut buf = match self.remain_bytes {
            0 => return Poll::Ready(None),
            n if n > SNAP_CHUNK_LEN => vec![0; SNAP_CHUNK_LEN],
            n => vec![0; n],
        };
        let result = self.snap.read_exact(buf.as_mut_slice());
        match result {
            Ok(_) => {
                self.remain_bytes -= buf.len();
                let mut chunk = SnapshotChunk::default();
                chunk.set_data(buf);
                Poll::Ready(Some(Ok((chunk, WriteFlags::default().buffer_hint(true)))))
            }
            Err(e) => Poll::Ready(Some(Err(box_err!("failed to read snapshot chunk: {}", e)))),
        }
    }
}

struct SendStat {
    key: SnapKey,
    total_size: u64,
    elapsed: Duration,
}

/// Send the snapshot to specified address.
///
/// It will first send the normal raft snapshot message and then send the snapshot file.
fn send_snap(
    env: Arc<Environment>,
    mgr: SnapManager,
    security_mgr: Arc<SecurityManager>,
    cfg: &Config,
    addr: &str,
    msg: RaftMessage,
) -> Result<impl Future<Output = Result<SendStat>>> {
    assert!(msg.get_message().has_snapshot());
    let timer = Instant::now();

    let send_timer = SEND_SNAP_HISTOGRAM.start_coarse_timer();

    let key = {
        let snap = msg.get_message().get_snapshot();
        SnapKey::from_snap(snap)?
    };

    mgr.register(key.clone(), SnapEntry::Sending);
    let deregister = {
        let (mgr, key) = (mgr.clone(), key.clone());
        DeferContext::new(move || mgr.deregister(&key, &SnapEntry::Sending))
    };

    let s = box_try!(mgr.get_snapshot_for_sending(&key));
    if !s.exists() {
        return Err(box_err!("missing snap file: {:?}", s.path()));
    }
    let total_size = s.total_size()?;

    let mut chunks = {
        let mut first_chunk = SnapshotChunk::default();
        first_chunk.set_message(msg);

        SnapChunk {
            first: Some(first_chunk),
            snap: s,
            remain_bytes: total_size as usize,
        }
    };

    let cb = ChannelBuilder::new(env)
        .stream_initial_window_size(cfg.grpc_stream_initial_window_size.0 as i32)
        .keepalive_time(cfg.grpc_keepalive_time.0)
        .keepalive_timeout(cfg.grpc_keepalive_timeout.0)
        .default_compression_algorithm(cfg.grpc_compression_algorithm());

    let channel = security_mgr.connect(cb, addr);
    let client = TikvClient::new(channel);
    let (sink, receiver) = client.snapshot()?;

    let send_task = async move {
        let mut sink = sink.sink_map_err(Error::from);
        sink.send_all(&mut chunks).await?;
        sink.close().await?;
        let recv_result = receiver.map_err(Error::from).await;
        send_timer.observe_duration();
        drop(deregister);
        drop(client);
        match recv_result {
            Ok(_) => {
                fail_point!("snapshot_delete_after_send");
                chunks.snap.delete();
                // TODO: improve it after rustc resolves the bug.
                // Call `info` in the closure directly will cause rustc
                // panic with `Cannot create local mono-item for DefId`.
                Ok(SendStat {
                    key,
                    total_size,
                    elapsed: timer.elapsed(),
                })
            }
            Err(e) => Err(e),
        }
    };
    Ok(send_task)
}

struct RecvSnapContext {
    key: SnapKey,
    file: Option<Box<dyn GenericSnapshot>>,
    raft_msg: RaftMessage,
    _with_io_type: WithIOType,
}

impl RecvSnapContext {
    fn new(head_chunk: Option<SnapshotChunk>, snap_mgr: &SnapManager) -> Result<Self> {
        // head_chunk is None means the stream is empty.
        let mut head = head_chunk.ok_or_else(|| Error::Other("empty gRPC stream".into()))?;
        if !head.has_message() {
            return Err(box_err!("no raft message in the first chunk"));
        }

        let meta = head.take_message();
        let key = match SnapKey::from_snap(meta.get_message().get_snapshot()) {
            Ok(k) => k,
            Err(e) => return Err(box_err!("failed to create snap key: {:?}", e)),
        };

        let data = meta.get_message().get_snapshot().get_data();
        let mut snapshot = RaftSnapshotData::default();
        snapshot.merge_from_bytes(data)?;
        let with_io_type = WithIOType::new(if snapshot.get_meta().get_for_balance() {
            IOType::LoadBalance
        } else {
            IOType::Replication
        });

        let snap = {
            let s = match snap_mgr.get_snapshot_for_receiving(&key, data) {
                Ok(s) => s,
                Err(e) => return Err(box_err!("{} failed to create snapshot file: {:?}", key, e)),
            };

            if s.exists() {
                let p = s.path();
                info!("snapshot file already exists, skip receiving"; "snap_key" => %key, "file" => p);
                None
            } else {
                Some(s)
            }
        };

        Ok(RecvSnapContext {
            key,
            file: snap,
            raft_msg: meta,
            _with_io_type: with_io_type,
        })
    }

    fn finish<R: RaftStoreRouter<impl KvEngine>>(self, raft_router: R) -> Result<()> {
        let key = self.key;
        if let Some(mut file) = self.file {
            info!("saving snapshot file"; "snap_key" => %key, "file" => file.path());
            if let Err(e) = file.save() {
                let path = file.path();
                let e = box_err!("{} failed to save snapshot file {}: {:?}", key, path, e);
                return Err(e);
            }
        }
        if let Err(e) = raft_router.send_raft_msg(self.raft_msg) {
            return Err(box_err!("{} failed to send snapshot to raft: {}", key, e));
        }
        Ok(())
    }
}

fn recv_snap<R: RaftStoreRouter<impl KvEngine> + 'static>(
    stream: RequestStream<SnapshotChunk>,
    sink: ClientStreamingSink<Done>,
    snap_mgr: SnapManager,
    raft_router: R,
) -> impl Future<Output = Result<()>> {
    let recv_task = async move {
        let mut stream = stream.map_err(Error::from);
        let head = stream.next().await.transpose()?;
        let mut context = RecvSnapContext::new(head, &snap_mgr)?;
        if context.file.is_none() {
            return context.finish(raft_router);
        }
        let context_key = context.key.clone();
        snap_mgr.register(context.key.clone(), SnapEntry::Receiving);

        while let Some(item) = stream.next().await {
            let mut chunk = item?;
            let data = chunk.take_data();
            if data.is_empty() {
                return Err(box_err!("{} receive chunk with empty data", context.key));
            }
            if let Err(e) = context.file.as_mut().unwrap().write_all(&data) {
                let key = &context.key;
                let path = context.file.as_mut().unwrap().path();
                let e = box_err!("{} failed to write snapshot file {}: {}", key, path, e);
                return Err(e);
            }
        }

        let res = context.finish(raft_router);
        snap_mgr.deregister(&context_key, &SnapEntry::Receiving);
        res
    };

    async move {
        match recv_task.await {
            Ok(()) => sink.success(Done::default()).await.map_err(Error::from),
            Err(e) => {
                let status = RpcStatus::new(RpcStatusCode::UNKNOWN, Some(format!("{:?}", e)));
                sink.fail(status).await.map_err(Error::from)
            }
        }
    }
}

pub struct Runner<E, R>
where
    E: KvEngine,
    R: RaftStoreRouter<E> + 'static,
{
    env: Arc<Environment>,
    snap_mgr: SnapManager,
    pool: Runtime,
    raft_router: R,
    security_mgr: Arc<SecurityManager>,
    cfg: Arc<Config>,
    sending_count: Arc<AtomicUsize>,
    recving_count: Arc<AtomicUsize>,
    engine: PhantomData<E>,
}

impl<E, R> Runner<E, R>
where
    E: KvEngine,
    R: RaftStoreRouter<E> + 'static,
{
    pub fn new(
        env: Arc<Environment>,
        snap_mgr: SnapManager,
        r: R,
        security_mgr: Arc<SecurityManager>,
        cfg: Arc<Config>,
    ) -> Runner<E, R> {
        Runner {
            env,
            snap_mgr,
            pool: RuntimeBuilder::new()
                .threaded_scheduler()
                .thread_name(thd_name!("snap-sender"))
                .core_threads(DEFAULT_POOL_SIZE)
                .on_thread_start(tikv_alloc::add_thread_memory_accessor)
                .on_thread_stop(tikv_alloc::remove_thread_memory_accessor)
                .build()
                .unwrap(),
            raft_router: r,
            security_mgr,
            cfg,
            sending_count: Arc::new(AtomicUsize::new(0)),
            recving_count: Arc::new(AtomicUsize::new(0)),
            engine: PhantomData,
        }
    }
}

impl<E, R> Runnable for Runner<E, R>
where
    E: KvEngine,
    R: RaftStoreRouter<E> + 'static,
{
    type Task = Task;

    fn run(&mut self, task: Task) {
        match task {
            Task::Recv { stream, sink } => {
                let task_num = self.recving_count.load(Ordering::SeqCst);
                if task_num >= self.cfg.concurrent_recv_snap_limit {
                    warn!("too many recving snapshot tasks, ignore");
                    let status = RpcStatus::new(
                        RpcStatusCode::RESOURCE_EXHAUSTED,
                        Some(format!(
                            "the number of received snapshot tasks {} exceeded the limitation {}",
                            task_num, self.cfg.concurrent_recv_snap_limit
                        )),
                    );
                    self.pool.spawn(sink.fail(status));
                    return;
                }
                SNAP_TASK_COUNTER_STATIC.recv.inc();

                let snap_mgr = self.snap_mgr.clone();
                let raft_router = self.raft_router.clone();
                let recving_count = Arc::clone(&self.recving_count);
                recving_count.fetch_add(1, Ordering::SeqCst);
                let task = async move {
                    let result = recv_snap(stream, sink, snap_mgr, raft_router).await;
                    recving_count.fetch_sub(1, Ordering::SeqCst);
                    if let Err(e) = result {
                        error!("failed to recv snapshot"; "err" => %e);
                    }
                };
                self.pool.spawn(task);
            }
            Task::Send { addr, msg, cb } => {
                fail_point!("send_snapshot");
                if self.sending_count.load(Ordering::SeqCst) >= self.cfg.concurrent_send_snap_limit
                {
                    warn!(
                        "too many sending snapshot tasks, drop Send Snap[to: {}, snap: {:?}]",
                        addr, msg
                    );
                    cb(Err(Error::Other("Too many sending snapshot tasks".into())));
                    return;
                }
                SNAP_TASK_COUNTER_STATIC.send.inc();

                let env = Arc::clone(&self.env);
                let mgr = self.snap_mgr.clone();
                let security_mgr = Arc::clone(&self.security_mgr);
                let sending_count = Arc::clone(&self.sending_count);
                sending_count.fetch_add(1, Ordering::SeqCst);

                let send_task = send_snap(env, mgr, security_mgr, &self.cfg, &addr, msg);
                let task = async move {
                    let res = match send_task {
                        Err(e) => Err(e),
                        Ok(f) => f.await,
                    };
                    match res {
                        Ok(stat) => {
                            info!(
                                "sent snapshot";
                                "region_id" => stat.key.region_id,
                                "snap_key" => %stat.key,
                                "size" => stat.total_size,
                                "duration" => ?stat.elapsed
                            );
                            cb(Ok(()));
                        }
                        Err(e) => {
                            error!("failed to send snap"; "to_addr" => addr, "err" => ?e);
                            cb(Err(e));
                        }
                    };
                    sending_count.fetch_sub(1, Ordering::SeqCst);
                };

                self.pool.spawn(task);
            }
        }
    }
}
<|endoftext|>


<commit_before>use core::ops::Mul;
use core::num::Wrapping;
use {One, CheckedMul};

/// Binary operator for raising a value to a power.
pub trait Pow<RHS> {
    /// The result after applying the operator.
    type Output;

    /// Returns `self` to the power `rhs`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::Pow;
    /// assert_eq!(Pow::pow(10u32, 2u32), 100);
    /// ```
    fn pow(self, rhs: RHS) -> Self::Output;
}

macro_rules! pow_impl {
    ($t:ty) => {
        pow_impl!($t, u8);
        pow_impl!($t, u16);
        pow_impl!($t, u32);
        pow_impl!($t, u64);
        pow_impl!($t, usize);
    };
    ($t:ty, $rhs:ty) => {
        pow_impl!($t, $rhs, |x, p| pow(x, p as usize));
    };
    ($t:ty, $rhs:ty, $method:expr) => {
        impl Pow<$rhs> for $t {
            type Output = $t;
            #[inline]
            fn pow(self, rhs: $rhs) -> $t {
                ($method)(self, rhs)
            }
        }

        impl<'a> Pow<&'a $rhs> for $t {
            type Output = $t;
            #[inline]
            fn pow(self, rhs: &'a $rhs) -> $t {
                ($method)(self, *rhs)
            }
        }

        impl<'a> Pow<$rhs> for &'a $t {
            type Output = $t;
            #[inline]
            fn pow(self, rhs: $rhs) -> $t {
                ($method)(*self, rhs)
            }
        }

        impl<'a, 'b> Pow<&'a $rhs> for &'b $t {
            type Output = $t;
            #[inline]
            fn pow(self, rhs: &'a $rhs) -> $t {
                ($method)(*self, *rhs)
            }
        }
    };
}

pow_impl!(u8, u8, |x: u8, p| x.pow(p as u32));
pow_impl!(u8, u16, |x: u8, p| x.pow(p as u32));
pow_impl!(u8, u32, u8::pow);
pow_impl!(u8, usize);
pow_impl!(i8, u8, |x: i8, p| x.pow(p as u32));
pow_impl!(i8, u16, |x: i8, p| x.pow(p as u32));
pow_impl!(i8, u32, i8::pow);
pow_impl!(i8, usize);
pow_impl!(u16, u8, |x: u16, p| x.pow(p as u32));
pow_impl!(u16, u16, |x: u16, p| x.pow(p as u32));
pow_impl!(u16, u32, u16::pow);
pow_impl!(u16, usize);
pow_impl!(i16, u8, |x: i16, p| x.pow(p as u32));
pow_impl!(i16, u16, |x: i16, p| x.pow(p as u32));
pow_impl!(i16, u32, i16::pow);
pow_impl!(i16, usize);
pow_impl!(u32, u8, |x: u32, p| x.pow(p as u32));
pow_impl!(u32, u16, |x: u32, p| x.pow(p as u32));
pow_impl!(u32, u32, u32::pow);
pow_impl!(u32, usize);
pow_impl!(i32, u8, |x: i32, p| x.pow(p as u32));
pow_impl!(i32, u16, |x: i32, p| x.pow(p as u32));
pow_impl!(i32, u32, i32::pow);
pow_impl!(i32, usize);
pow_impl!(u64, u8, |x: u64, p| x.pow(p as u32));
pow_impl!(u64, u16, |x: u64, p| x.pow(p as u32));
pow_impl!(u64, u32, u64::pow);
pow_impl!(u64, usize);
pow_impl!(i64, u8, |x: i64, p| x.pow(p as u32));
pow_impl!(i64, u16, |x: i64, p| x.pow(p as u32));
pow_impl!(i64, u32, i64::pow);
pow_impl!(i64, usize);
pow_impl!(usize, u8, |x: usize, p| x.pow(p as u32));
pow_impl!(usize, u16, |x: usize, p| x.pow(p as u32));
pow_impl!(usize, u32, usize::pow);
pow_impl!(usize, usize);
pow_impl!(isize, u8, |x: isize, p| x.pow(p as u32));
pow_impl!(isize, u16, |x: isize, p| x.pow(p as u32));
pow_impl!(isize, u32, isize::pow);
pow_impl!(isize, usize);
pow_impl!(Wrapping<u8>);
pow_impl!(Wrapping<i8>);
pow_impl!(Wrapping<u16>);
pow_impl!(Wrapping<i16>);
pow_impl!(Wrapping<u32>);
pow_impl!(Wrapping<i32>);
pow_impl!(Wrapping<u64>);
pow_impl!(Wrapping<i64>);
pow_impl!(Wrapping<usize>);
pow_impl!(Wrapping<isize>);

#[cfg(feature = "std")]
mod float_impls {
    use super::Pow;

    pow_impl!(f32, i8, |x: f32, p| x.powi(p as i32));
    pow_impl!(f32, u8, |x: f32, p| x.powi(p as i32));
    pow_impl!(f32, i16, |x: f32, p| x.powi(p as i32));
    pow_impl!(f32, u16, |x: f32, p| x.powi(p as i32));
    pow_impl!(f32, i32, f32::powi);
    pow_impl!(f64, i8, |x: f64, p| x.powi(p as i32));
    pow_impl!(f64, u8, |x: f64, p| x.powi(p as i32));
    pow_impl!(f64, i16, |x: f64, p| x.powi(p as i32));
    pow_impl!(f64, u16, |x: f64, p| x.powi(p as i32));
    pow_impl!(f64, i32, f64::powi);
    pow_impl!(f32, f32, f32::powf);
    pow_impl!(f64, f32, |x: f64, p| x.powf(p as f64));
    pow_impl!(f64, f64, f64::powf);
}

/// Raises a value to the power of exp, using exponentiation by squaring.
///
/// # Example
///
/// ```rust
/// use num_traits::pow;
///
/// assert_eq!(pow(2i8, 4), 16);
/// assert_eq!(pow(6u8, 3), 216);
/// ```
#[inline]
pub fn pow<T: Clone + One + Mul<T, Output = T>>(mut base: T, mut exp: usize) -> T {
    if exp == 0 { return T::one() }

    while exp & 1 == 0 {
        base = base.clone() * base;
        exp >>= 1;
    }
    if exp == 1 { return base }

    let mut acc = base.clone();
    while exp > 1 {
        exp >>= 1;
        base = base.clone() * base;
        if exp & 1 == 1 {
            acc = acc * base.clone();
        }
    }
    acc
}

/// Raises a value to the power of exp, returning `None` if an overflow occurred.
///
/// Otherwise same as the `pow` function.
///
/// # Example
///
/// ```rust
/// use num_traits::checked_pow;
///
/// assert_eq!(checked_pow(2i8, 4), Some(16));
/// assert_eq!(checked_pow(7i8, 8), None);
/// assert_eq!(checked_pow(7u32, 8), Some(5_764_801));
/// ```
#[inline]
pub fn checked_pow<T: Clone + One + CheckedMul>(mut base: T, mut exp: usize) -> Option<T> {
    if exp == 0 { return Some(T::one()) }

    macro_rules! optry {
        ( $ expr : expr ) => {
            if let Some(val) = $expr { val } else { return None }
        }
    }

    while exp & 1 == 0 {
        base = optry!(base.checked_mul(&base));
        exp >>= 1;
    }
    if exp == 1 { return Some(base) }

    let mut acc = base.clone();
    while exp > 1 {
        exp >>= 1;
        base = optry!(base.checked_mul(&base));
        if exp & 1 == 1 {
            acc = optry!(acc.checked_mul(&base));
        }
    }
    Some(acc)
}
<commit_msg>Remove Pow<u64> (accidentally added).<commit_after>use core::ops::Mul;
use core::num::Wrapping;
use {One, CheckedMul};

/// Binary operator for raising a value to a power.
pub trait Pow<RHS> {
    /// The result after applying the operator.
    type Output;

    /// Returns `self` to the power `rhs`.
    ///
    /// # Examples
    ///
    /// ```
    /// use num_traits::Pow;
    /// assert_eq!(Pow::pow(10u32, 2u32), 100);
    /// ```
    fn pow(self, rhs: RHS) -> Self::Output;
}

macro_rules! pow_impl {
    ($t:ty) => {
        pow_impl!($t, u8);
        pow_impl!($t, u16);
        pow_impl!($t, u32);
        pow_impl!($t, usize);
    };
    ($t:ty, $rhs:ty) => {
        pow_impl!($t, $rhs, |x, p| pow(x, p as usize));
    };
    ($t:ty, $rhs:ty, $method:expr) => {
        impl Pow<$rhs> for $t {
            type Output = $t;
            #[inline]
            fn pow(self, rhs: $rhs) -> $t {
                ($method)(self, rhs)
            }
        }

        impl<'a> Pow<&'a $rhs> for $t {
            type Output = $t;
            #[inline]
            fn pow(self, rhs: &'a $rhs) -> $t {
                ($method)(self, *rhs)
            }
        }

        impl<'a> Pow<$rhs> for &'a $t {
            type Output = $t;
            #[inline]
            fn pow(self, rhs: $rhs) -> $t {
                ($method)(*self, rhs)
            }
        }

        impl<'a, 'b> Pow<&'a $rhs> for &'b $t {
            type Output = $t;
            #[inline]
            fn pow(self, rhs: &'a $rhs) -> $t {
                ($method)(*self, *rhs)
            }
        }
    };
}

pow_impl!(u8, u8, |x: u8, p| x.pow(p as u32));
pow_impl!(u8, u16, |x: u8, p| x.pow(p as u32));
pow_impl!(u8, u32, u8::pow);
pow_impl!(u8, usize);
pow_impl!(i8, u8, |x: i8, p| x.pow(p as u32));
pow_impl!(i8, u16, |x: i8, p| x.pow(p as u32));
pow_impl!(i8, u32, i8::pow);
pow_impl!(i8, usize);
pow_impl!(u16, u8, |x: u16, p| x.pow(p as u32));
pow_impl!(u16, u16, |x: u16, p| x.pow(p as u32));
pow_impl!(u16, u32, u16::pow);
pow_impl!(u16, usize);
pow_impl!(i16, u8, |x: i16, p| x.pow(p as u32));
pow_impl!(i16, u16, |x: i16, p| x.pow(p as u32));
pow_impl!(i16, u32, i16::pow);
pow_impl!(i16, usize);
pow_impl!(u32, u8, |x: u32, p| x.pow(p as u32));
pow_impl!(u32, u16, |x: u32, p| x.pow(p as u32));
pow_impl!(u32, u32, u32::pow);
pow_impl!(u32, usize);
pow_impl!(i32, u8, |x: i32, p| x.pow(p as u32));
pow_impl!(i32, u16, |x: i32, p| x.pow(p as u32));
pow_impl!(i32, u32, i32::pow);
pow_impl!(i32, usize);
pow_impl!(u64, u8, |x: u64, p| x.pow(p as u32));
pow_impl!(u64, u16, |x: u64, p| x.pow(p as u32));
pow_impl!(u64, u32, u64::pow);
pow_impl!(u64, usize);
pow_impl!(i64, u8, |x: i64, p| x.pow(p as u32));
pow_impl!(i64, u16, |x: i64, p| x.pow(p as u32));
pow_impl!(i64, u32, i64::pow);
pow_impl!(i64, usize);
pow_impl!(usize, u8, |x: usize, p| x.pow(p as u32));
pow_impl!(usize, u16, |x: usize, p| x.pow(p as u32));
pow_impl!(usize, u32, usize::pow);
pow_impl!(usize, usize);
pow_impl!(isize, u8, |x: isize, p| x.pow(p as u32));
pow_impl!(isize, u16, |x: isize, p| x.pow(p as u32));
pow_impl!(isize, u32, isize::pow);
pow_impl!(isize, usize);
pow_impl!(Wrapping<u8>);
pow_impl!(Wrapping<i8>);
pow_impl!(Wrapping<u16>);
pow_impl!(Wrapping<i16>);
pow_impl!(Wrapping<u32>);
pow_impl!(Wrapping<i32>);
pow_impl!(Wrapping<u64>);
pow_impl!(Wrapping<i64>);
pow_impl!(Wrapping<usize>);
pow_impl!(Wrapping<isize>);

#[cfg(feature = "std")]
mod float_impls {
    use super::Pow;

    pow_impl!(f32, i8, |x: f32, p| x.powi(p as i32));
    pow_impl!(f32, u8, |x: f32, p| x.powi(p as i32));
    pow_impl!(f32, i16, |x: f32, p| x.powi(p as i32));
    pow_impl!(f32, u16, |x: f32, p| x.powi(p as i32));
    pow_impl!(f32, i32, f32::powi);
    pow_impl!(f64, i8, |x: f64, p| x.powi(p as i32));
    pow_impl!(f64, u8, |x: f64, p| x.powi(p as i32));
    pow_impl!(f64, i16, |x: f64, p| x.powi(p as i32));
    pow_impl!(f64, u16, |x: f64, p| x.powi(p as i32));
    pow_impl!(f64, i32, f64::powi);
    pow_impl!(f32, f32, f32::powf);
    pow_impl!(f64, f32, |x: f64, p| x.powf(p as f64));
    pow_impl!(f64, f64, f64::powf);
}

/// Raises a value to the power of exp, using exponentiation by squaring.
///
/// # Example
///
/// ```rust
/// use num_traits::pow;
///
/// assert_eq!(pow(2i8, 4), 16);
/// assert_eq!(pow(6u8, 3), 216);
/// ```
#[inline]
pub fn pow<T: Clone + One + Mul<T, Output = T>>(mut base: T, mut exp: usize) -> T {
    if exp == 0 { return T::one() }

    while exp & 1 == 0 {
        base = base.clone() * base;
        exp >>= 1;
    }
    if exp == 1 { return base }

    let mut acc = base.clone();
    while exp > 1 {
        exp >>= 1;
        base = base.clone() * base;
        if exp & 1 == 1 {
            acc = acc * base.clone();
        }
    }
    acc
}

/// Raises a value to the power of exp, returning `None` if an overflow occurred.
///
/// Otherwise same as the `pow` function.
///
/// # Example
///
/// ```rust
/// use num_traits::checked_pow;
///
/// assert_eq!(checked_pow(2i8, 4), Some(16));
/// assert_eq!(checked_pow(7i8, 8), None);
/// assert_eq!(checked_pow(7u32, 8), Some(5_764_801));
/// ```
#[inline]
pub fn checked_pow<T: Clone + One + CheckedMul>(mut base: T, mut exp: usize) -> Option<T> {
    if exp == 0 { return Some(T::one()) }

    macro_rules! optry {
        ( $ expr : expr ) => {
            if let Some(val) = $expr { val } else { return None }
        }
    }

    while exp & 1 == 0 {
        base = optry!(base.checked_mul(&base));
        exp >>= 1;
    }
    if exp == 1 { return Some(base) }

    let mut acc = base.clone();
    while exp > 1 {
        exp >>= 1;
        base = optry!(base.checked_mul(&base));
        if exp & 1 == 1 {
            acc = optry!(acc.checked_mul(&base));
        }
    }
    Some(acc)
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::fmt;

use serde::ser::{Serialize, Serializer, SerializeStruct};
use serde::de::{self, Deserialize, Deserializer, DeserializeSeed, IntoDeserializer, MapAccess, Unexpected, Visitor};
use serde::de::value::BorrowedStrDeserializer;

use error::Error;

/// Represents any valid JSON value as a series of raw bytes.
///
/// This type can be used to defer parsing parts of a payload until later,
/// or to embed it verbatim into another JSON payload.
///
/// When serializing, a value of this type will retain its original formatting
/// and will not be minified or pretty-printed.
///
/// When deserializing, this type can not be used with the `#[serde(flatten)]` attribute,
/// as it relies on the original input buffer.

#[derive(Debug, Clone)]
pub struct RawValue<'a>(Cow<'a, str>);

impl<'a> AsRef<str> for RawValue<'a> {
    fn as_ref(&self) -> &str {
        &self.0
    }
}

impl<'a> fmt::Display for RawValue<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(&self.0)
    }
}

pub const SERDE_STRUCT_NAME: &'static str = "$serde_json::RawValue";
pub const SERDE_STRUCT_FIELD_NAME: &'static str = "$serde_json::RawValue::id";

impl<'a> Serialize for RawValue<'a> {
    #[inline]
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut s = serializer.serialize_struct(SERDE_STRUCT_NAME, 1)?;
        s.serialize_field(SERDE_STRUCT_FIELD_NAME, &self.0)?;
        s.end()
    }
}

impl<'a, 'de: 'a> Deserialize<'de> for RawValue<'a> {
    #[inline]
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct RawValueVisitor;

        impl<'de> Visitor<'de> for RawValueVisitor {
            type Value = RawValue<'de>;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                write!(formatter, "any valid JSON value")
            }

            fn visit_map<V>(self, mut visitor: V) -> Result<Self::Value, V::Error>
            where
                V: MapAccess<'de>,
            {
                let value = visitor.next_key::<RawKey>()?;
                if value.is_none() {
                    return Err(de::Error::invalid_type(Unexpected::Map, &self));
                }
                visitor.next_value_seed(RawValueFromString)
            }
        }

        deserializer.deserialize_newtype_struct(SERDE_STRUCT_NAME, RawValueVisitor)
    }
}

struct RawKey;

impl<'de> Deserialize<'de> for RawKey {
    fn deserialize<D>(deserializer: D) -> Result<RawKey, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct FieldVisitor;

        impl<'de> Visitor<'de> for FieldVisitor {
            type Value = ();

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("raw value")
            }

            fn visit_str<E>(self, s: &str) -> Result<(), E>
            where
                E: de::Error,
            {
                if s == SERDE_STRUCT_FIELD_NAME {
                    Ok(())
                } else {
                    Err(de::Error::custom("unexpected raw value"))
                }
            }
        }

        deserializer.deserialize_identifier(FieldVisitor)?;
        Ok(RawKey)
    }
}

pub struct RawValueFromString;

impl<'de> DeserializeSeed<'de> for RawValueFromString {
    type Value = RawValue<'de>;

    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_str(self)
    }
}

impl<'de> Visitor<'de> for RawValueFromString {
    type Value = RawValue<'de>;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("raw value")
    }

    fn visit_borrowed_str<E>(self, s: &'de str) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(RawValue(Cow::Borrowed(s)))
    }

    fn visit_string<E>(self, s: String) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(RawValue(Cow::Owned(s)))
    }
}

struct RawKeyDeserializer;

impl<'de> Deserializer<'de> for RawKeyDeserializer {
    type Error = Error;

    fn deserialize_any<V>(self, visitor: V) -> Result<V::Value, Error>
    where
        V: de::Visitor<'de>,
    {
        visitor.visit_borrowed_str(SERDE_STRUCT_FIELD_NAME)
    }

    forward_to_deserialize_any! {
        bool u8 u16 u32 u64 u128 i8 i16 i32 i64 i128 f32 f64 char str string seq
        bytes byte_buf map struct option unit newtype_struct ignored_any
        unit_struct tuple_struct tuple enum identifier
    }
}

pub struct OwnedRawDeserializer {
    pub raw_value: Option<String>,
}

impl<'de> MapAccess<'de> for OwnedRawDeserializer {
    type Error = Error;

    fn next_key_seed<K>(&mut self, seed: K) -> Result<Option<K::Value>, Error>
    where
        K: de::DeserializeSeed<'de>,
    {
        if self.raw_value.is_none() {
            return Ok(None);
        }
        seed.deserialize(RawKeyDeserializer).map(Some)
    }

    fn next_value_seed<V>(&mut self, seed: V) -> Result<V::Value, Error>
    where
        V: de::DeserializeSeed<'de>,
    {
        seed.deserialize(self.raw_value.take().unwrap().into_deserializer())
    }
}

pub struct BorrowedRawDeserializer<'de> {
    pub raw_value: Option<&'de str>,
}

impl<'de> MapAccess<'de> for BorrowedRawDeserializer<'de> {
    type Error = Error;

    fn next_key_seed<K>(&mut self, seed: K) -> Result<Option<K::Value>, Error>
    where
        K: de::DeserializeSeed<'de>,
    {
        if self.raw_value.is_none() {
            return Ok(None);
        }
        seed.deserialize(RawKeyDeserializer).map(Some)
    }

    fn next_value_seed<V>(&mut self, seed: V) -> Result<V::Value, Error>
    where
        V: de::DeserializeSeed<'de>,
    {
        seed.deserialize(BorrowedStrDeserializer::new(self.raw_value.take().unwrap()))
    }
}
<commit_msg>Handwrite Debug for RawValue to minimize escaping<commit_after>use std::borrow::Cow;
use std::fmt::{self, Debug};

use serde::ser::{Serialize, Serializer, SerializeStruct};
use serde::de::{self, Deserialize, Deserializer, DeserializeSeed, IntoDeserializer, MapAccess, Unexpected, Visitor};
use serde::de::value::BorrowedStrDeserializer;

use error::Error;

/// Represents any valid JSON value as a series of raw bytes.
///
/// This type can be used to defer parsing parts of a payload until later,
/// or to embed it verbatim into another JSON payload.
///
/// When serializing, a value of this type will retain its original formatting
/// and will not be minified or pretty-printed.
///
/// When deserializing, this type can not be used with the `#[serde(flatten)]` attribute,
/// as it relies on the original input buffer.

#[derive(Clone)]
pub struct RawValue<'a>(Cow<'a, str>);

impl<'a> Debug for RawValue<'a> {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter
            .debug_tuple("RawValue")
            .field(&format_args!("{}", self.0))
            .finish()
    }
}

impl<'a> AsRef<str> for RawValue<'a> {
    fn as_ref(&self) -> &str {
        &self.0
    }
}

impl<'a> fmt::Display for RawValue<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(&self.0)
    }
}

pub const SERDE_STRUCT_NAME: &'static str = "$serde_json::RawValue";
pub const SERDE_STRUCT_FIELD_NAME: &'static str = "$serde_json::RawValue::id";

impl<'a> Serialize for RawValue<'a> {
    #[inline]
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut s = serializer.serialize_struct(SERDE_STRUCT_NAME, 1)?;
        s.serialize_field(SERDE_STRUCT_FIELD_NAME, &self.0)?;
        s.end()
    }
}

impl<'a, 'de: 'a> Deserialize<'de> for RawValue<'a> {
    #[inline]
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct RawValueVisitor;

        impl<'de> Visitor<'de> for RawValueVisitor {
            type Value = RawValue<'de>;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                write!(formatter, "any valid JSON value")
            }

            fn visit_map<V>(self, mut visitor: V) -> Result<Self::Value, V::Error>
            where
                V: MapAccess<'de>,
            {
                let value = visitor.next_key::<RawKey>()?;
                if value.is_none() {
                    return Err(de::Error::invalid_type(Unexpected::Map, &self));
                }
                visitor.next_value_seed(RawValueFromString)
            }
        }

        deserializer.deserialize_newtype_struct(SERDE_STRUCT_NAME, RawValueVisitor)
    }
}

struct RawKey;

impl<'de> Deserialize<'de> for RawKey {
    fn deserialize<D>(deserializer: D) -> Result<RawKey, D::Error>
    where
        D: Deserializer<'de>,
    {
        struct FieldVisitor;

        impl<'de> Visitor<'de> for FieldVisitor {
            type Value = ();

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("raw value")
            }

            fn visit_str<E>(self, s: &str) -> Result<(), E>
            where
                E: de::Error,
            {
                if s == SERDE_STRUCT_FIELD_NAME {
                    Ok(())
                } else {
                    Err(de::Error::custom("unexpected raw value"))
                }
            }
        }

        deserializer.deserialize_identifier(FieldVisitor)?;
        Ok(RawKey)
    }
}

pub struct RawValueFromString;

impl<'de> DeserializeSeed<'de> for RawValueFromString {
    type Value = RawValue<'de>;

    fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_str(self)
    }
}

impl<'de> Visitor<'de> for RawValueFromString {
    type Value = RawValue<'de>;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("raw value")
    }

    fn visit_borrowed_str<E>(self, s: &'de str) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(RawValue(Cow::Borrowed(s)))
    }

    fn visit_string<E>(self, s: String) -> Result<Self::Value, E>
    where
        E: de::Error,
    {
        Ok(RawValue(Cow::Owned(s)))
    }
}

struct RawKeyDeserializer;

impl<'de> Deserializer<'de> for RawKeyDeserializer {
    type Error = Error;

    fn deserialize_any<V>(self, visitor: V) -> Result<V::Value, Error>
    where
        V: de::Visitor<'de>,
    {
        visitor.visit_borrowed_str(SERDE_STRUCT_FIELD_NAME)
    }

    forward_to_deserialize_any! {
        bool u8 u16 u32 u64 u128 i8 i16 i32 i64 i128 f32 f64 char str string seq
        bytes byte_buf map struct option unit newtype_struct ignored_any
        unit_struct tuple_struct tuple enum identifier
    }
}

pub struct OwnedRawDeserializer {
    pub raw_value: Option<String>,
}

impl<'de> MapAccess<'de> for OwnedRawDeserializer {
    type Error = Error;

    fn next_key_seed<K>(&mut self, seed: K) -> Result<Option<K::Value>, Error>
    where
        K: de::DeserializeSeed<'de>,
    {
        if self.raw_value.is_none() {
            return Ok(None);
        }
        seed.deserialize(RawKeyDeserializer).map(Some)
    }

    fn next_value_seed<V>(&mut self, seed: V) -> Result<V::Value, Error>
    where
        V: de::DeserializeSeed<'de>,
    {
        seed.deserialize(self.raw_value.take().unwrap().into_deserializer())
    }
}

pub struct BorrowedRawDeserializer<'de> {
    pub raw_value: Option<&'de str>,
}

impl<'de> MapAccess<'de> for BorrowedRawDeserializer<'de> {
    type Error = Error;

    fn next_key_seed<K>(&mut self, seed: K) -> Result<Option<K::Value>, Error>
    where
        K: de::DeserializeSeed<'de>,
    {
        if self.raw_value.is_none() {
            return Ok(None);
        }
        seed.deserialize(RawKeyDeserializer).map(Some)
    }

    fn next_value_seed<V>(&mut self, seed: V) -> Result<V::Value, Error>
    where
        V: de::DeserializeSeed<'de>,
    {
        seed.deserialize(BorrowedStrDeserializer::new(self.raw_value.take().unwrap()))
    }
}
<|endoftext|>


<commit_before>use Meta;
use cgmath::{Matrix4, Point3};
use collision::{Aabb3, Frustum, Relation};
use iron;
use iron::mime::Mime;
use iron::prelude::*;
use quadtree::{ChildIndex, Node};
use router::Router;
use std::fs;
use std::io::Read;
use std::path::PathBuf;
use std::sync::Arc;
use urlencoded::UrlEncodedQuery;

#[derive(Serialize, Debug)]
struct BoundingRect {
    min_x: f32,
    min_y: f32,
    edge_length: f32,
}

#[derive(Serialize, Debug)]
struct NodeMeta {
    id: String,
    bounding_rect: BoundingRect,
}

#[derive(Serialize, Debug)]
struct MetaReply {
    bounding_rect: BoundingRect,
    tile_size: u32,
    deepest_level: u8,
}

struct HandleNodeImage {
    directory: PathBuf,
}

impl iron::Handler for HandleNodeImage {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let id = req.extensions.get::<Router>().unwrap().find("id");
        if id.is_none() {
            return Ok(Response::with(iron::status::NotFound));
        }
        let id = id.unwrap();

        let mut filename = self.directory.join(id);
        filename.set_extension("png");
        let mut file = itry!(fs::File::open(filename), iron::status::NotFound);

        let mut reply = Vec::new();
        itry!(file.read_to_end(&mut reply));
        let content_type = "image/png".parse::<Mime>().unwrap();
        Ok(Response::with((content_type, iron::status::Ok, reply)))
    }
}

struct HandleMeta {
    meta: Arc<Meta>,
}

impl iron::Handler for HandleMeta {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        let result = MetaReply {
            bounding_rect: BoundingRect {
                min_x: self.meta.bounding_rect.min().x,
                min_y: self.meta.bounding_rect.min().y,
                edge_length: self.meta.bounding_rect.edge_length(),
            },
            tile_size: self.meta.tile_size,
            deepest_level: self.meta.deepest_level,
        };
        let reply = ::serde_json::to_string_pretty(&result).unwrap();
        let content_type = "application/json".parse::<Mime>().unwrap();
        Ok(Response::with((content_type, iron::status::Ok, reply)))
    }
}

struct HandleNodesForLevel {
    meta: Arc<Meta>,
}

impl iron::Handler for HandleNodesForLevel {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let query = req.get_ref::<UrlEncodedQuery>().unwrap();
        let level_s = &query.get("level").unwrap()[0];
        let level = itry!(level_s.parse::<u8>(), iron::status::NotFound);

        let matrix = {
            // Entries are column major.
            let e: Vec<f32> = query.get("matrix").unwrap()[0]
                .split(',')
                .map(|s| s.parse::<f32>().unwrap())
                .collect();
            Matrix4::new(
                e[0],
                e[1],
                e[2],
                e[3],
                e[4],
                e[5],
                e[6],
                e[7],
                e[8],
                e[9],
                e[10],
                e[11],
                e[12],
                e[13],
                e[14],
                e[15],
            )
        };
        let frustum = Frustum::from_matrix4(matrix).unwrap();

        let mut result = Vec::new();
        let mut open = vec![
            Node::root_with_bounding_rect(self.meta.bounding_rect.clone()),
        ];
        while !open.is_empty() {
            let node = open.pop().unwrap();
            if node.level() == level {
                let aabb = Aabb3::new(
                    Point3::new(node.bounding_rect.min().x, node.bounding_rect.min().y, -0.1),
                    Point3::new(node.bounding_rect.max().x, node.bounding_rect.max().y, 0.1),
                );
                if frustum.contains(&aabb) == Relation::Out || !self.meta.nodes.contains(&node.id) {
                    continue;
                }
                result.push(NodeMeta {
                    id: node.id.to_string(),
                    bounding_rect: BoundingRect {
                        min_x: node.bounding_rect.min().x,
                        min_y: node.bounding_rect.min().y,
                        edge_length: node.bounding_rect.edge_length(),
                    },
                });
            } else {
                for i in 0..4 {
                    open.push(node.get_child(ChildIndex::from_u8(i)));
                }
            }
        }

        let reply = ::serde_json::to_string_pretty(&result).unwrap();
        let content_type = "application/json".parse::<Mime>().unwrap();
        Ok(Response::with((content_type, iron::status::Ok, reply)))
    }
}

pub fn serve(prefix: &str, router: &mut Router, quadtree_directory: PathBuf) {
    let meta = Arc::new(Meta::from_disk(quadtree_directory.join("meta.pb")));
    router.get(
        format!("{}/meta", prefix),
        HandleMeta {
            meta: Arc::clone(&meta),
        },
    );
    router.get(
        format!("{}/nodes_for_level", prefix),
        HandleNodesForLevel {
            meta: Arc::clone(&meta),
        },
    );
    router.get(
        format!("{}/node_image/:id", prefix),
        HandleNodeImage {
            directory: quadtree_directory.clone(),
        },
    );
}
<commit_msg>Actually cull the quadtree while traversing. (#171)<commit_after>use Meta;
use cgmath::{Matrix4, Point3};
use collision::{Aabb3, Frustum, Relation};
use iron;
use iron::mime::Mime;
use iron::prelude::*;
use quadtree::{ChildIndex, Node};
use router::Router;
use std::fs;
use std::io::Read;
use std::path::PathBuf;
use std::sync::Arc;
use urlencoded::UrlEncodedQuery;

#[derive(Serialize, Debug)]
struct BoundingRect {
    min_x: f32,
    min_y: f32,
    edge_length: f32,
}

#[derive(Serialize, Debug)]
struct NodeMeta {
    id: String,
    bounding_rect: BoundingRect,
}

#[derive(Serialize, Debug)]
struct MetaReply {
    bounding_rect: BoundingRect,
    tile_size: u32,
    deepest_level: u8,
}

struct HandleNodeImage {
    directory: PathBuf,
}

impl iron::Handler for HandleNodeImage {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let id = req.extensions.get::<Router>().unwrap().find("id");
        if id.is_none() {
            return Ok(Response::with(iron::status::NotFound));
        }
        let id = id.unwrap();

        let mut filename = self.directory.join(id);
        filename.set_extension("png");
        let mut file = itry!(fs::File::open(filename), iron::status::NotFound);

        let mut reply = Vec::new();
        itry!(file.read_to_end(&mut reply));
        let content_type = "image/png".parse::<Mime>().unwrap();
        Ok(Response::with((content_type, iron::status::Ok, reply)))
    }
}

struct HandleMeta {
    meta: Arc<Meta>,
}

impl iron::Handler for HandleMeta {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        let result = MetaReply {
            bounding_rect: BoundingRect {
                min_x: self.meta.bounding_rect.min().x,
                min_y: self.meta.bounding_rect.min().y,
                edge_length: self.meta.bounding_rect.edge_length(),
            },
            tile_size: self.meta.tile_size,
            deepest_level: self.meta.deepest_level,
        };
        let reply = ::serde_json::to_string_pretty(&result).unwrap();
        let content_type = "application/json".parse::<Mime>().unwrap();
        Ok(Response::with((content_type, iron::status::Ok, reply)))
    }
}

struct HandleNodesForLevel {
    meta: Arc<Meta>,
}

impl iron::Handler for HandleNodesForLevel {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        let query = req.get_ref::<UrlEncodedQuery>().unwrap();
        let level_s = &query.get("level").unwrap()[0];
        let level = itry!(level_s.parse::<u8>(), iron::status::NotFound);

        let matrix = {
            // Entries are column major.
            let e: Vec<f32> = query.get("matrix").unwrap()[0]
                .split(',')
                .map(|s| s.parse::<f32>().unwrap())
                .collect();
            Matrix4::new(
                e[0],
                e[1],
                e[2],
                e[3],
                e[4],
                e[5],
                e[6],
                e[7],
                e[8],
                e[9],
                e[10],
                e[11],
                e[12],
                e[13],
                e[14],
                e[15],
            )
        };
        let frustum = Frustum::from_matrix4(matrix).unwrap();

        // TODO(sirver): This function could actually work much faster by not traversing the
        // levels, but just finding the covering of the rectangle of the current bounding box.
        //
        // Also it should probably not take a frustum but the view bounding box we are interested
        // in.
        let mut result = Vec::new();
        let mut open = vec![
            Node::root_with_bounding_rect(self.meta.bounding_rect.clone()),
        ];
        while !open.is_empty() {
            let node = open.pop().unwrap();
            let aabb = Aabb3::new(
                Point3::new(node.bounding_rect.min().x, node.bounding_rect.min().y, -0.1),
                Point3::new(node.bounding_rect.max().x, node.bounding_rect.max().y, 0.1),
            );

            if frustum.contains(&aabb) == Relation::Out || !self.meta.nodes.contains(&node.id) {
                continue;
            }

            if node.level() == level {
                result.push(NodeMeta {
                    id: node.id.to_string(),
                    bounding_rect: BoundingRect {
                        min_x: node.bounding_rect.min().x,
                        min_y: node.bounding_rect.min().y,
                        edge_length: node.bounding_rect.edge_length(),
                    },
                });
            } else {
                for i in 0..4 {
                    open.push(node.get_child(ChildIndex::from_u8(i)));
                }
            }
        }

        let reply = ::serde_json::to_string_pretty(&result).unwrap();
        let content_type = "application/json".parse::<Mime>().unwrap();
        Ok(Response::with((content_type, iron::status::Ok, reply)))
    }
}

pub fn serve(prefix: &str, router: &mut Router, quadtree_directory: PathBuf) {
    let meta = Arc::new(Meta::from_disk(quadtree_directory.join("meta.pb")));
    router.get(
        format!("{}/meta", prefix),
        HandleMeta {
            meta: Arc::clone(&meta),
        },
    );
    router.get(
        format!("{}/nodes_for_level", prefix),
        HandleNodesForLevel {
            meta: Arc::clone(&meta),
        },
    );
    router.get(
        format!("{}/node_image/:id", prefix),
        HandleNodeImage {
            directory: quadtree_directory.clone(),
        },
    );
}
<|endoftext|>


<commit_before>use cli::{clear_screen, Args};
use env_logger;
use error::Result;
use gitignore;
use log;
use notification_filter::NotificationFilter;
#[cfg(target_os = "linux")]
use notify;
use pathop::PathOp;
use process::{self, Process};
use signal::{self, Signal};
use std::{
    collections::HashMap,
    io::Write,
    sync::{
        mpsc::{channel, Receiver},
        Arc, RwLock,
    },
    time::Duration,
};
use watcher::{Event, Watcher};

fn init_logger(debug: bool) {
    let mut log_builder = env_logger::Builder::new();
    let level = if debug {
        log::LevelFilter::Debug
    } else {
        log::LevelFilter::Warn
    };

    log_builder
        .format(|buf, r| writeln!(buf, "*** {}", r.args()))
        .filter(None, level)
        .init();
}

pub trait Handler {
    /// Initialises the `Handler` with a copy of the arguments.
    fn new(args: Args) -> Result<Self>
    where
        Self: Sized;

    /// Called through a manual request, such as an initial run.
    ///
    /// # Returns
    ///
    /// A `Result` which means:
    ///
    /// - `Err`: an error has occurred while processing, quit.
    /// - `Ok(false)`: everything is fine and the loop can continue.
    /// - `Ok(true)`: everything is fine but we should gracefully stop.
    fn on_manual(&mut self) -> Result<bool>;

    /// Called through a file-update request.
    ///
    /// # Parameters
    ///
    /// - `ops`: The list of events that triggered this update.
    ///
    /// # Returns
    ///
    /// A `Result` which means:
    ///
    /// - `Err`: an error has occurred while processing, quit.
    /// - `Ok(true)`: everything is fine and the loop can continue.
    /// - `Ok(false)`: everything is fine but we should gracefully stop.
    fn on_update(&mut self, ops: Vec<PathOp>) -> Result<bool>;
}

/// Starts watching, and calls a handler when something happens.
///
/// Given an argument structure and a `Handler` type, starts the watcher
/// loop (blocking until done).
pub fn watch<H>(args: Args) -> Result<()>
where
    H: Handler,
{
    init_logger(args.debug);
    let mut handler = H::new(args.clone())?;

    let gitignore = gitignore::load(if args.no_vcs_ignore { &[] } else { &args.paths });
    let filter = NotificationFilter::new(&args.filters, &args.ignores, gitignore)?;

    let (tx, rx) = channel();
    let poll = args.poll.clone();
    #[cfg(target_os = "linux")]
    let poll_interval = args.poll_interval.clone();
    let watcher = Watcher::new(tx.clone(), &args.paths, args.poll, args.poll_interval).or_else(|err| {
        if poll {
            return Err(err);
        }

        #[cfg(target_os = "linux")]
        {
            use nix::libc;
            let mut fallback = false;
            if let notify::Error::Io(ref e) = err {
                if e.raw_os_error() == Some(libc::ENOSPC) {
                    warn!("System notification limit is too small, falling back to polling mode. For better performance increase system limit:\n\tsysctl fs.inotify.max_user_watches=524288");
                    fallback = true;
                }
            }

            if fallback {
                return Watcher::new(tx, &args.paths, true, poll_interval);
            }
        }

        Err(err)
    })?;

    if watcher.is_polling() {
        warn!("Polling for changes every {} ms", args.poll_interval);
    }

    // Call handler initially, if necessary
    if args.run_initially && !args.once {
        if args.clear_screen {
            clear_screen();
        }

        if !handler.on_manual()? {
            return Ok(());
        }
    }

    loop {
        debug!("Waiting for filesystem activity");
        let paths = wait_fs(&rx, &filter, args.debounce);
        debug!("Paths updated: {:?}", paths);

        if args.clear_screen {
            clear_screen();
        }

        if !handler.on_update(paths)? {
            break;
        }
    }

    Ok(())
}

pub struct ExecHandler {
    args: Args,
    signal: Option<Signal>,
    child_process: Arc<RwLock<Option<Process>>>,
}

impl Handler for ExecHandler {
    fn new(args: Args) -> Result<Self> {
        let child_process: Arc<RwLock<Option<Process>>> = Arc::new(RwLock::new(None));
        let weak_child = Arc::downgrade(&child_process);

        // Convert signal string to the corresponding integer
        let signal = signal::new(args.signal.clone());

        signal::install_handler(move |sig: Signal| {
            if let Some(lock) = weak_child.upgrade() {
                let strong = lock.read().unwrap();
                if let Some(ref child) = *strong {
                    match sig {
                        Signal::SIGCHLD => child.reap(), // SIGCHLD is special, initiate reap()
                        _ => child.signal(sig),
                    }
                }
            }
        });

        Ok(Self {
            args,
            signal,
            child_process,
        })
    }

    fn on_manual(&mut self) -> Result<bool> {
        let mut guard = self.child_process.write()?;
        *guard = Some(process::spawn(
            &self.args.cmd,
            Vec::new(),
            self.args.no_shell,
        ));

        Ok(true)
    }

    fn on_update(&mut self, ops: Vec<PathOp>) -> Result<bool> {
        // We have three scenarios here:
        //
        // 1. Make sure the previous run was ended, then run the command again
        // 2. Just send a specified signal to the child, do nothing more
        // 3. Send SIGTERM to the child, wait for it to exit, then run the command again
        // 4. Send a specified signal to the child, wait for it to exit, then run the command again
        //
        let scenario = (self.args.restart, self.signal.is_some());

        match scenario {
            // Custom restart behaviour (--restart was given, and --signal specified):
            // Send specified signal to the child, wait for it to exit, then run the command again
            (true, true) => {
                signal_process(&self.child_process, self.signal, true);

                // Launch child process
                if self.args.clear_screen {
                    clear_screen();
                }

                debug!("Launching child process");
                {
                    let mut guard = self.child_process.write()?;
                    *guard = Some(process::spawn(&self.args.cmd, ops, self.args.no_shell));
                }
            }

            // Default restart behaviour (--restart was given, but --signal wasn't specified):
            // Send SIGTERM to the child, wait for it to exit, then run the command again
            (true, false) => {
                let sigterm = signal::new(Some("SIGTERM".into()));
                signal_process(&self.child_process, sigterm, true);

                // Launch child process
                if self.args.clear_screen {
                    clear_screen();
                }

                debug!("Launching child process");
                {
                    let mut guard = self.child_process.write()?;
                    *guard = Some(process::spawn(&self.args.cmd, ops, self.args.no_shell));
                }
            }

            // SIGHUP scenario: --signal was given, but --restart was not
            // Just send a signal (e.g. SIGHUP) to the child, do nothing more
            (false, true) => signal_process(&self.child_process, self.signal, false),

            // Default behaviour (neither --signal nor --restart specified):
            // Make sure the previous run was ended, then run the command again
            (false, false) => {
                signal_process(&self.child_process, None, true);

                // Launch child process
                if self.args.clear_screen {
                    clear_screen();
                }

                debug!("Launching child process");
                {
                    let mut guard = self.child_process.write()?;
                    *guard = Some(process::spawn(&self.args.cmd, ops, self.args.no_shell));
                }
            }
        }

        // Handle once option for integration testing
        if self.args.once {
            signal_process(&self.child_process, self.signal, false);
            return Ok(false);
        }

        Ok(true)
    }
}

pub fn run(args: Args) -> Result<()> {
    watch::<ExecHandler>(args)
}

fn wait_fs(rx: &Receiver<Event>, filter: &NotificationFilter, debounce: u64) -> Vec<PathOp> {
    let mut paths = vec![];
    let mut cache = HashMap::new();

    loop {
        let e = rx.recv().expect("error when reading event");

        if let Some(ref path) = e.path {
            let pathop = PathOp::new(path, e.op.ok(), e.cookie);
            // Ignore cache for the initial file. Otherwise, in
            // debug mode it's hard to track what's going on
            let excluded = filter.is_excluded(path);
            if !cache.contains_key(&pathop) {
                cache.insert(pathop.clone(), excluded);
            }

            if !excluded {
                paths.push(pathop);
                break;
            }
        }
    }

    // Wait for filesystem activity to cool off
    let timeout = Duration::from_millis(debounce);
    while let Ok(e) = rx.recv_timeout(timeout) {
        if let Some(ref path) = e.path {
            let pathop = PathOp::new(path, e.op.ok(), e.cookie);
            if cache.contains_key(&pathop) {
                continue;
            }

            let excluded = filter.is_excluded(path);

            cache.insert(pathop.clone(), excluded);

            if !excluded {
                paths.push(pathop);
            }
        }
    }

    paths
}

// signal_process sends signal to process. It waits for the process to exit if wait is true
fn signal_process(process: &RwLock<Option<Process>>, signal: Option<Signal>, wait: bool) {
    let guard = process.read().unwrap();

    if let Some(ref child) = *guard {
        if let Some(s) = signal {
            child.signal(s);
        }

        if wait {
            debug!("Waiting for process to exit...");
            child.wait();
        }
    }
}
<commit_msg>[run] Dry spawning the child process<commit_after>use cli::{clear_screen, Args};
use env_logger;
use error::Result;
use gitignore;
use log;
use notification_filter::NotificationFilter;
#[cfg(target_os = "linux")]
use notify;
use pathop::PathOp;
use process::{self, Process};
use signal::{self, Signal};
use std::{
    collections::HashMap,
    io::Write,
    sync::{
        mpsc::{channel, Receiver},
        Arc, RwLock,
    },
    time::Duration,
};
use watcher::{Event, Watcher};

fn init_logger(debug: bool) {
    let mut log_builder = env_logger::Builder::new();
    let level = if debug {
        log::LevelFilter::Debug
    } else {
        log::LevelFilter::Warn
    };

    log_builder
        .format(|buf, r| writeln!(buf, "*** {}", r.args()))
        .filter(None, level)
        .init();
}

pub trait Handler {
    /// Initialises the `Handler` with a copy of the arguments.
    fn new(args: Args) -> Result<Self>
    where
        Self: Sized;

    /// Called through a manual request, such as an initial run.
    ///
    /// # Returns
    ///
    /// A `Result` which means:
    ///
    /// - `Err`: an error has occurred while processing, quit.
    /// - `Ok(false)`: everything is fine and the loop can continue.
    /// - `Ok(true)`: everything is fine but we should gracefully stop.
    fn on_manual(&mut self) -> Result<bool>;

    /// Called through a file-update request.
    ///
    /// # Parameters
    ///
    /// - `ops`: The list of events that triggered this update.
    ///
    /// # Returns
    ///
    /// A `Result` which means:
    ///
    /// - `Err`: an error has occurred while processing, quit.
    /// - `Ok(true)`: everything is fine and the loop can continue.
    /// - `Ok(false)`: everything is fine but we should gracefully stop.
    fn on_update(&mut self, ops: Vec<PathOp>) -> Result<bool>;
}

/// Starts watching, and calls a handler when something happens.
///
/// Given an argument structure and a `Handler` type, starts the watcher
/// loop (blocking until done).
pub fn watch<H>(args: Args) -> Result<()>
where
    H: Handler,
{
    init_logger(args.debug);
    let mut handler = H::new(args.clone())?;

    let gitignore = gitignore::load(if args.no_vcs_ignore { &[] } else { &args.paths });
    let filter = NotificationFilter::new(&args.filters, &args.ignores, gitignore)?;

    let (tx, rx) = channel();
    let poll = args.poll.clone();
    #[cfg(target_os = "linux")]
    let poll_interval = args.poll_interval.clone();
    let watcher = Watcher::new(tx.clone(), &args.paths, args.poll, args.poll_interval).or_else(|err| {
        if poll {
            return Err(err);
        }

        #[cfg(target_os = "linux")]
        {
            use nix::libc;
            let mut fallback = false;
            if let notify::Error::Io(ref e) = err {
                if e.raw_os_error() == Some(libc::ENOSPC) {
                    warn!("System notification limit is too small, falling back to polling mode. For better performance increase system limit:\n\tsysctl fs.inotify.max_user_watches=524288");
                    fallback = true;
                }
            }

            if fallback {
                return Watcher::new(tx, &args.paths, true, poll_interval);
            }
        }

        Err(err)
    })?;

    if watcher.is_polling() {
        warn!("Polling for changes every {} ms", args.poll_interval);
    }

    // Call handler initially, if necessary
    if args.run_initially && !args.once {
        if args.clear_screen {
            clear_screen();
        }

        if !handler.on_manual()? {
            return Ok(());
        }
    }

    loop {
        debug!("Waiting for filesystem activity");
        let paths = wait_fs(&rx, &filter, args.debounce);
        debug!("Paths updated: {:?}", paths);

        if args.clear_screen {
            clear_screen();
        }

        if !handler.on_update(paths)? {
            break;
        }
    }

    Ok(())
}

pub struct ExecHandler {
    args: Args,
    signal: Option<Signal>,
    child_process: Arc<RwLock<Option<Process>>>,
}

impl ExecHandler {
    fn spawn(&mut self, ops: Vec<PathOp>) -> Result<()> {
        let mut guard = self.child_process.write()?;
        *guard = Some(process::spawn(
            &self.args.cmd,
            ops,
            self.args.no_shell,
        ));

        Ok(())
    }
}

impl Handler for ExecHandler {
    fn new(args: Args) -> Result<Self> {
        let child_process: Arc<RwLock<Option<Process>>> = Arc::new(RwLock::new(None));
        let weak_child = Arc::downgrade(&child_process);

        // Convert signal string to the corresponding integer
        let signal = signal::new(args.signal.clone());

        signal::install_handler(move |sig: Signal| {
            if let Some(lock) = weak_child.upgrade() {
                let strong = lock.read().unwrap();
                if let Some(ref child) = *strong {
                    match sig {
                        Signal::SIGCHLD => child.reap(), // SIGCHLD is special, initiate reap()
                        _ => child.signal(sig),
                    }
                }
            }
        });

        Ok(Self {
            args,
            signal,
            child_process,
        })
    }

    // Only returns Err() on lock poisoning.
    fn on_manual(&mut self) -> Result<bool> {
        self.spawn(Vec::new()).and(Ok(true))
    }

    // Only returns Err() on lock poisoning.
    fn on_update(&mut self, ops: Vec<PathOp>) -> Result<bool> {
        // We have three scenarios here:
        //
        // 1. Make sure the previous run was ended, then run the command again
        // 2. Just send a specified signal to the child, do nothing more
        // 3. Send SIGTERM to the child, wait for it to exit, then run the command again
        // 4. Send a specified signal to the child, wait for it to exit, then run the command again
        //
        let scenario = (self.args.restart, self.signal.is_some());

        match scenario {
            // Custom restart behaviour (--restart was given, and --signal specified):
            // Send specified signal to the child, wait for it to exit, then run the command again
            (true, true) => {
                signal_process(&self.child_process, self.signal, true);

                // Launch child process
                if self.args.clear_screen {
                    clear_screen();
                }

                debug!("Launching child process");
                self.spawn(ops)?;
            }

            // Default restart behaviour (--restart was given, but --signal wasn't specified):
            // Send SIGTERM to the child, wait for it to exit, then run the command again
            (true, false) => {
                let sigterm = signal::new(Some("SIGTERM".into()));
                signal_process(&self.child_process, sigterm, true);

                // Launch child process
                if self.args.clear_screen {
                    clear_screen();
                }

                debug!("Launching child process");
                self.spawn(ops)?;
            }

            // SIGHUP scenario: --signal was given, but --restart was not
            // Just send a signal (e.g. SIGHUP) to the child, do nothing more
            (false, true) => signal_process(&self.child_process, self.signal, false),

            // Default behaviour (neither --signal nor --restart specified):
            // Make sure the previous run was ended, then run the command again
            (false, false) => {
                signal_process(&self.child_process, None, true);

                // Launch child process
                if self.args.clear_screen {
                    clear_screen();
                }

                debug!("Launching child process");
                self.spawn(ops)?;
            }
        }

        // Handle once option for integration testing
        if self.args.once {
            signal_process(&self.child_process, self.signal, false);
            return Ok(false);
        }

        Ok(true)
    }
}

pub fn run(args: Args) -> Result<()> {
    watch::<ExecHandler>(args)
}

fn wait_fs(rx: &Receiver<Event>, filter: &NotificationFilter, debounce: u64) -> Vec<PathOp> {
    let mut paths = vec![];
    let mut cache = HashMap::new();

    loop {
        let e = rx.recv().expect("error when reading event");

        if let Some(ref path) = e.path {
            let pathop = PathOp::new(path, e.op.ok(), e.cookie);
            // Ignore cache for the initial file. Otherwise, in
            // debug mode it's hard to track what's going on
            let excluded = filter.is_excluded(path);
            if !cache.contains_key(&pathop) {
                cache.insert(pathop.clone(), excluded);
            }

            if !excluded {
                paths.push(pathop);
                break;
            }
        }
    }

    // Wait for filesystem activity to cool off
    let timeout = Duration::from_millis(debounce);
    while let Ok(e) = rx.recv_timeout(timeout) {
        if let Some(ref path) = e.path {
            let pathop = PathOp::new(path, e.op.ok(), e.cookie);
            if cache.contains_key(&pathop) {
                continue;
            }

            let excluded = filter.is_excluded(path);

            cache.insert(pathop.clone(), excluded);

            if !excluded {
                paths.push(pathop);
            }
        }
    }

    paths
}

// signal_process sends signal to process. It waits for the process to exit if wait is true
fn signal_process(process: &RwLock<Option<Process>>, signal: Option<Signal>, wait: bool) {
    let guard = process.read().unwrap();

    if let Some(ref child) = *guard {
        if let Some(s) = signal {
            child.signal(s);
        }

        if wait {
            debug!("Waiting for process to exit...");
            child.wait();
        }
    }
}
<|endoftext|>