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

use std::sync::{atomic::AtomicUsize, atomic::Ordering, Arc};
use std::time::{Duration, Instant};

use futures::{
    channel::mpsc::{
        channel as bounded, unbounded, Receiver, SendError as FuturesSendError, Sender,
        TrySendError, UnboundedReceiver, UnboundedSender,
    },
    executor::block_on,
    stream, SinkExt, Stream, StreamExt,
};
use grpcio::WriteFlags;
use kvproto::cdcpb::ChangeDataEvent;

use tikv_util::{impl_display_as_debug, warn};

use crate::service::{CdcEvent, EventBatcher};

const CDC_MSG_MAX_BATCH_SIZE: usize = 128;
// Assume the average size of event is 1KB.
// 2 = (CDC_MSG_MAX_BATCH_SIZE * 1KB / service::CDC_MAX_RESP_SIZE).ceil() + 1 /* reserve for ResolvedTs */;
pub const CDC_EVENT_MAX_BATCH_SIZE: usize = 2;

#[derive(Clone)]
pub struct MemoryQuota {
    capacity: usize,
    in_use: Arc<AtomicUsize>,
}

impl MemoryQuota {
    pub fn new(capacity: usize) -> MemoryQuota {
        MemoryQuota {
            capacity,
            in_use: Arc::new(AtomicUsize::new(0)),
        }
    }
    pub fn in_use(&self) -> usize {
        self.in_use.load(Ordering::Relaxed)
    }
    pub fn cap(&self) -> usize {
        self.capacity
    }
    fn alloc(&self, bytes: usize) -> bool {
        let mut in_use_bytes = self.in_use.load(Ordering::Relaxed);
        loop {
            if in_use_bytes + bytes > self.capacity {
                return false;
            }
            let new_in_use_bytes = in_use_bytes + bytes;
            match self.in_use.compare_exchange(
                in_use_bytes,
                new_in_use_bytes,
                Ordering::Acquire,
                Ordering::Relaxed,
            ) {
                Ok(_) => return true,
                Err(current) => in_use_bytes = current,
            }
        }
    }
    fn free(&self, bytes: usize) {
        self.in_use.fetch_sub(bytes, Ordering::Release);
    }
}

pub fn channel(buffer: usize, memory_quota: MemoryQuota) -> (Sink, Drain) {
    let (unbounded_sender, unbounded_receiver) = unbounded();
    let (bounded_sender, bounded_receiver) = bounded(buffer);
    (
        Sink {
            unbounded_sender,
            bounded_sender,
            memory_quota: memory_quota.clone(),
        },
        Drain {
            unbounded_receiver,
            bounded_receiver,
            memory_quota,
        },
    )
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SendError {
    Full,
    Disconnected,
    Congested,
}

impl std::error::Error for SendError {}

impl_display_as_debug!(SendError);

macro_rules! impl_from_future_send_error {
    ($($f:ty,)+) => {
        $(
            impl From<$f> for SendError {
                fn from(e: $f) -> Self {
                    if e.is_disconnected() {
                        SendError::Disconnected
                    } else if e.is_full() {
                        Self::Full
                    } else {
                        unreachable!()
                    }
                }
            }
        )+
    };
}

impl_from_future_send_error! {
    FuturesSendError,
    TrySendError<(CdcEvent, usize)>,
}

#[derive(Clone)]
pub struct Sink {
    unbounded_sender: UnboundedSender<(CdcEvent, usize)>,
    bounded_sender: Sender<(CdcEvent, usize)>,
    memory_quota: MemoryQuota,
}

impl Sink {
    pub fn unbounded_send(&self, event: CdcEvent, force: bool) -> Result<(), SendError> {
        // Try it's best to send error events.
        let bytes = if !force { event.size() as usize } else { 0 };
        if bytes != 0 && !self.memory_quota.alloc(bytes) {
            return Err(SendError::Congested);
        }
        self.unbounded_sender
            .unbounded_send((event, bytes))
            .map_err(SendError::from)
    }

    pub async fn send_all(&mut self, events: Vec<CdcEvent>) -> Result<(), SendError> {
        for event in events {
            let bytes = event.size() as usize;
            if !self.memory_quota.alloc(bytes) {
                return Err(SendError::Congested);
            }
            self.bounded_sender.feed((event, bytes)).await?;
        }
        self.bounded_sender.flush().await?;
        Ok(())
    }
}

pub struct Drain {
    unbounded_receiver: UnboundedReceiver<(CdcEvent, usize)>,
    bounded_receiver: Receiver<(CdcEvent, usize)>,
    memory_quota: MemoryQuota,
}

impl<'a> Drain {
    pub fn drain(&'a mut self) -> impl Stream<Item = (CdcEvent, usize)> + 'a {
        stream::select(&mut self.bounded_receiver, &mut self.unbounded_receiver).map(
            |(mut event, size)| {
                if let CdcEvent::Barrier(ref mut barrier) = event {
                    if let Some(barrier) = barrier.take() {
                        // Unset barrier when it is received.
                        barrier(());
                    }
                }
                (event, size)
            },
        )
    }

    // Forwards contents to the sink, simulates StreamExt::forward.
    pub async fn forward<S, E>(&'a mut self, sink: &mut S) -> Result<(), E>
    where
        S: futures::Sink<(ChangeDataEvent, WriteFlags), Error = E> + Unpin,
    {
        let memory_quota = self.memory_quota.clone();
        let mut chunks = self.drain().ready_chunks(CDC_MSG_MAX_BATCH_SIZE);
        while let Some(events) = chunks.next().await {
            let mut bytes = 0;
            let mut batcher = EventBatcher::with_capacity(CDC_EVENT_MAX_BATCH_SIZE);
            events.into_iter().for_each(|(e, size)| {
                bytes += size;
                batcher.push(e);
            });
            let resps = batcher.build();
            let last_idx = resps.len() - 1;
            // Events are about to be sent, free pending events memory counter.
            memory_quota.free(bytes as _);
            for (i, e) in resps.into_iter().enumerate() {
                // Buffer messages and flush them at once.
                let write_flags = WriteFlags::default().buffer_hint(i != last_idx);
                sink.feed((e, write_flags)).await?;
            }
            sink.flush().await?;
        }
        Ok(())
    }
}

impl Drop for Drain {
    fn drop(&mut self) {
        self.bounded_receiver.close();
        self.unbounded_receiver.close();
        let start = Instant::now();
        let mut drain = Box::pin(async {
            let mut total_bytes = 0;
            while let Some((_, bytes)) = self.drain().next().await {
                total_bytes += bytes;
            }
            self.memory_quota.free(total_bytes);
        });
        block_on(&mut drain);
        let takes = start.elapsed();
        if takes >= Duration::from_millis(200) {
            warn!("drop Drain too slow"; "takes" => ?takes);
        }
    }
}

#[cfg(test)]
pub fn recv_timeout<S, I>(s: &mut S, dur: std::time::Duration) -> Result<Option<I>, ()>
where
    S: Stream<Item = I> + Unpin,
{
    poll_timeout(&mut s.next(), dur)
}

#[cfg(test)]
pub fn poll_timeout<F, I>(fut: &mut F, dur: std::time::Duration) -> Result<I, ()>
where
    F: std::future::Future<Output = I> + Unpin,
{
    use futures::FutureExt;
    let mut timeout = futures_timer::Delay::new(dur).fuse();
    let mut f = fut.fuse();
    futures::executor::block_on(async {
        futures::select! {
            () = timeout => Err(()),
            item = f => Ok(item),
        }
    })
}

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

    use std::sync::mpsc;
    use std::time::Duration;

    type Send = Box<dyn FnMut(CdcEvent) -> Result<(), SendError>>;
    fn new_test_cancal(buffer: usize, capacity: usize, force_send: bool) -> (Send, Drain) {
        let memory_quota = MemoryQuota::new(capacity);
        let (mut tx, rx) = channel(buffer, memory_quota);
        let mut flag = true;
        let send = move |event| {
            flag = !flag;
            if flag {
                tx.unbounded_send(event, force_send)
            } else {
                block_on(tx.send_all(vec![event]))
            }
        };
        (Box::new(send), rx)
    }

    #[test]
    fn test_barrier() {
        let force_send = false;
        let (mut send, mut rx) = new_test_cancal(10, usize::MAX, force_send);
        send(CdcEvent::Event(Default::default())).unwrap();
        let (btx1, brx1) = mpsc::channel();
        send(CdcEvent::Barrier(Some(Box::new(move |()| {
            btx1.send(()).unwrap();
        }))))
        .unwrap();
        send(CdcEvent::ResolvedTs(Default::default())).unwrap();
        let (btx2, brx2) = mpsc::channel();
        send(CdcEvent::Barrier(Some(Box::new(move |()| {
            btx2.send(()).unwrap();
        }))))
        .unwrap();

        let mut drain = rx.drain();
        brx1.recv_timeout(Duration::from_millis(100)).unwrap_err();
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
        assert_matches!(block_on(drain.next()), Some((CdcEvent::Event(_), _)));
        brx1.recv_timeout(Duration::from_millis(100)).unwrap_err();
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
        assert_matches!(block_on(drain.next()), Some((CdcEvent::Barrier(_), _)));
        brx1.recv_timeout(Duration::from_millis(100)).unwrap();
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
        assert_matches!(block_on(drain.next()), Some((CdcEvent::ResolvedTs(_), _)));
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
        assert_matches!(block_on(drain.next()), Some((CdcEvent::Barrier(_), _)));
        brx2.recv_timeout(Duration::from_millis(100)).unwrap();
        brx1.recv_timeout(Duration::from_millis(100)).unwrap_err();
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
    }

    #[test]
    fn test_nonblocking_batch() {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let force_send = false;
        let (mut send, mut drain) =
            new_test_cancal(CDC_MSG_MAX_BATCH_SIZE * 2, usize::MAX, force_send);
        let (mut tx, mut rx) = unbounded();
        runtime.spawn(async move {
            drain.forward(&mut tx).await.unwrap();
        });
        for count in 1..CDC_EVENT_MAX_BATCH_SIZE + CDC_EVENT_MAX_BATCH_SIZE / 2 {
            for _ in 0..count {
                send(CdcEvent::Event(Default::default())).unwrap();
            }
            recv_timeout(&mut rx, Duration::from_millis(100)).unwrap();
        }

        if recv_timeout(&mut rx, Duration::from_millis(100)).is_ok() {
            panic!("expect to be timeout");
        }
    }

    #[test]
    fn test_congest() {
        let mut e = kvproto::cdcpb::Event::default();
        e.region_id = 1;
        let event = CdcEvent::Event(e.clone());
        assert!(event.size() != 0);
        // 1KB
        let max_pending_bytes = 1024;
        let buffer = max_pending_bytes / event.size();
        let force_send = false;
        let (mut send, _rx) = new_test_cancal(buffer as _, max_pending_bytes as _, force_send);
        for _ in 0..buffer {
            send(CdcEvent::Event(e.clone())).unwrap();
        }
        assert_matches!(send(CdcEvent::Event(e)).unwrap_err(), SendError::Congested);
    }

    #[test]
    fn test_force_send() {
        let mut e = kvproto::cdcpb::Event::default();
        e.region_id = 1;
        let event = CdcEvent::Event(e.clone());
        assert!(event.size() != 0);
        // 1KB
        let max_pending_bytes = 1024;
        let buffer = max_pending_bytes / event.size();
        let memory_quota = MemoryQuota::new(max_pending_bytes as _);
        let (tx, _rx) = channel(buffer as _, memory_quota);
        for _ in 0..buffer {
            tx.unbounded_send(CdcEvent::Event(e.clone()), false)
                .unwrap();
        }
        assert_matches!(
            tx.unbounded_send(CdcEvent::Event(e.clone()), false)
                .unwrap_err(),
            SendError::Congested
        );
        tx.unbounded_send(CdcEvent::Event(e), true).unwrap();
    }

    #[test]
    fn test_channel_memory_leak() {
        let mut e = kvproto::cdcpb::Event::default();
        e.region_id = 1;
        let event = CdcEvent::Event(e.clone());
        assert!(event.size() != 0);
        // 1KB
        let max_pending_bytes = 1024;
        let buffer = max_pending_bytes / event.size() + 1;
        let force_send = false;
        // Make sure memory quota is freed when rx is dropped before tx.
        {
            let (mut send, rx) = new_test_cancal(buffer as _, max_pending_bytes as _, force_send);
            loop {
                match send(CdcEvent::Event(e.clone())) {
                    Ok(_) => (),
                    Err(e) => {
                        assert_matches!(e, SendError::Congested);
                        break;
                    }
                }
            }
            let memory_quota = rx.memory_quota.clone();
            assert_eq!(memory_quota.alloc(event.size() as _), false,);
            drop(rx);
            assert_eq!(memory_quota.alloc(1024), true,);
        }
        // Make sure memory quota is freed when tx is dropped before rx.
        {
            let (mut send, rx) = new_test_cancal(buffer as _, max_pending_bytes as _, force_send);
            loop {
                match send(CdcEvent::Event(e.clone())) {
                    Ok(_) => (),
                    Err(e) => {
                        assert_matches!(e, SendError::Congested);
                        break;
                    }
                }
            }
            let memory_quota = rx.memory_quota.clone();
            assert_eq!(memory_quota.alloc(event.size() as _), false,);
            drop(send);
            drop(rx);
            assert_eq!(memory_quota.alloc(1024), true,);
        }
    }
}
<commit_msg>cdc: fix sink quota leak on closed channel (#10413)<commit_after>// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.

use std::sync::{atomic::AtomicUsize, atomic::Ordering, Arc};
use std::time::{Duration, Instant};

use futures::{
    channel::mpsc::{
        channel as bounded, unbounded, Receiver, SendError as FuturesSendError, Sender,
        TrySendError, UnboundedReceiver, UnboundedSender,
    },
    executor::block_on,
    stream, SinkExt, Stream, StreamExt,
};
use grpcio::WriteFlags;
use kvproto::cdcpb::ChangeDataEvent;

use tikv_util::{impl_display_as_debug, warn};

use crate::service::{CdcEvent, EventBatcher};

const CDC_MSG_MAX_BATCH_SIZE: usize = 128;
// Assume the average size of event is 1KB.
// 2 = (CDC_MSG_MAX_BATCH_SIZE * 1KB / service::CDC_MAX_RESP_SIZE).ceil() + 1 /* reserve for ResolvedTs */;
pub const CDC_EVENT_MAX_BATCH_SIZE: usize = 2;

#[derive(Clone)]
pub struct MemoryQuota {
    capacity: usize,
    in_use: Arc<AtomicUsize>,
}

impl MemoryQuota {
    pub fn new(capacity: usize) -> MemoryQuota {
        MemoryQuota {
            capacity,
            in_use: Arc::new(AtomicUsize::new(0)),
        }
    }
    pub fn in_use(&self) -> usize {
        self.in_use.load(Ordering::Relaxed)
    }
    pub fn cap(&self) -> usize {
        self.capacity
    }
    fn alloc(&self, bytes: usize) -> bool {
        let mut in_use_bytes = self.in_use.load(Ordering::Relaxed);
        loop {
            if in_use_bytes + bytes > self.capacity {
                return false;
            }
            let new_in_use_bytes = in_use_bytes + bytes;
            match self.in_use.compare_exchange(
                in_use_bytes,
                new_in_use_bytes,
                Ordering::Acquire,
                Ordering::Relaxed,
            ) {
                Ok(_) => return true,
                Err(current) => in_use_bytes = current,
            }
        }
    }
    fn free(&self, bytes: usize) {
        self.in_use.fetch_sub(bytes, Ordering::Release);
    }
}

pub fn channel(buffer: usize, memory_quota: MemoryQuota) -> (Sink, Drain) {
    let (unbounded_sender, unbounded_receiver) = unbounded();
    let (bounded_sender, bounded_receiver) = bounded(buffer);
    (
        Sink {
            unbounded_sender,
            bounded_sender,
            memory_quota: memory_quota.clone(),
        },
        Drain {
            unbounded_receiver,
            bounded_receiver,
            memory_quota,
        },
    )
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SendError {
    Full,
    Disconnected,
    Congested,
}

impl std::error::Error for SendError {}

impl_display_as_debug!(SendError);

macro_rules! impl_from_future_send_error {
    ($($f:ty,)+) => {
        $(
            impl From<$f> for SendError {
                fn from(e: $f) -> Self {
                    if e.is_disconnected() {
                        SendError::Disconnected
                    } else if e.is_full() {
                        Self::Full
                    } else {
                        unreachable!()
                    }
                }
            }
        )+
    };
}

impl_from_future_send_error! {
    FuturesSendError,
    TrySendError<(CdcEvent, usize)>,
}

#[derive(Clone)]
pub struct Sink {
    unbounded_sender: UnboundedSender<(CdcEvent, usize)>,
    bounded_sender: Sender<(CdcEvent, usize)>,
    memory_quota: MemoryQuota,
}

impl Sink {
    pub fn unbounded_send(&self, event: CdcEvent, force: bool) -> Result<(), SendError> {
        // Try it's best to send error events.
        let bytes = if !force { event.size() as usize } else { 0 };
        if bytes != 0 && !self.memory_quota.alloc(bytes) {
            return Err(SendError::Congested);
        }
        match self.unbounded_sender.unbounded_send((event, bytes)) {
            Ok(_) => Ok(()),
            Err(e) => {
                // Free quota if send fails.
                self.memory_quota.free(bytes);
                Err(SendError::from(e))
            }
        }
    }

    pub async fn send_all(&mut self, events: Vec<CdcEvent>) -> Result<(), SendError> {
        // Allocate quota in advance.
        let mut total_bytes = 0;
        for event in &events {
            total_bytes += event.size();
        }
        if !self.memory_quota.alloc(total_bytes as _) {
            return Err(SendError::Congested);
        }
        for event in events {
            let bytes = event.size() as usize;
            if let Err(e) = self.bounded_sender.feed((event, bytes)).await {
                // Free quota if send fails.
                self.memory_quota.free(total_bytes as _);
                return Err(SendError::from(e));
            }
        }
        if let Err(e) = self.bounded_sender.flush().await {
            // Free quota if send fails.
            self.memory_quota.free(total_bytes as _);
            return Err(SendError::from(e));
        }
        Ok(())
    }
}

pub struct Drain {
    unbounded_receiver: UnboundedReceiver<(CdcEvent, usize)>,
    bounded_receiver: Receiver<(CdcEvent, usize)>,
    memory_quota: MemoryQuota,
}

impl<'a> Drain {
    pub fn drain(&'a mut self) -> impl Stream<Item = (CdcEvent, usize)> + 'a {
        stream::select(&mut self.bounded_receiver, &mut self.unbounded_receiver).map(
            |(mut event, size)| {
                if let CdcEvent::Barrier(ref mut barrier) = event {
                    if let Some(barrier) = barrier.take() {
                        // Unset barrier when it is received.
                        barrier(());
                    }
                }
                (event, size)
            },
        )
    }

    // Forwards contents to the sink, simulates StreamExt::forward.
    pub async fn forward<S, E>(&'a mut self, sink: &mut S) -> Result<(), E>
    where
        S: futures::Sink<(ChangeDataEvent, WriteFlags), Error = E> + Unpin,
    {
        let memory_quota = self.memory_quota.clone();
        let mut chunks = self.drain().ready_chunks(CDC_MSG_MAX_BATCH_SIZE);
        while let Some(events) = chunks.next().await {
            let mut bytes = 0;
            let mut batcher = EventBatcher::with_capacity(CDC_EVENT_MAX_BATCH_SIZE);
            events.into_iter().for_each(|(e, size)| {
                bytes += size;
                batcher.push(e);
            });
            let resps = batcher.build();
            let last_idx = resps.len() - 1;
            // Events are about to be sent, free pending events memory counter.
            memory_quota.free(bytes as _);
            for (i, e) in resps.into_iter().enumerate() {
                // Buffer messages and flush them at once.
                let write_flags = WriteFlags::default().buffer_hint(i != last_idx);
                sink.feed((e, write_flags)).await?;
            }
            sink.flush().await?;
        }
        Ok(())
    }
}

impl Drop for Drain {
    fn drop(&mut self) {
        self.bounded_receiver.close();
        self.unbounded_receiver.close();
        let start = Instant::now();
        let mut drain = Box::pin(async {
            let memory_quota = self.memory_quota.clone();
            let mut total_bytes = 0;
            let mut drain = self.drain();
            while let Some((_, bytes)) = drain.next().await {
                total_bytes += bytes;
            }
            memory_quota.free(total_bytes);
        });
        block_on(&mut drain);
        let takes = start.elapsed();
        if takes >= Duration::from_millis(200) {
            warn!("drop Drain too slow"; "takes" => ?takes);
        }
    }
}

#[cfg(test)]
pub fn recv_timeout<S, I>(s: &mut S, dur: std::time::Duration) -> Result<Option<I>, ()>
where
    S: Stream<Item = I> + Unpin,
{
    poll_timeout(&mut s.next(), dur)
}

#[cfg(test)]
pub fn poll_timeout<F, I>(fut: &mut F, dur: std::time::Duration) -> Result<I, ()>
where
    F: std::future::Future<Output = I> + Unpin,
{
    use futures::FutureExt;
    let mut timeout = futures_timer::Delay::new(dur).fuse();
    let mut f = fut.fuse();
    futures::executor::block_on(async {
        futures::select! {
            () = timeout => Err(()),
            item = f => Ok(item),
        }
    })
}

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

    use std::sync::mpsc;
    use std::time::Duration;

    type Send = Box<dyn FnMut(CdcEvent) -> Result<(), SendError>>;
    fn new_test_cancal(buffer: usize, capacity: usize, force_send: bool) -> (Send, Drain) {
        let memory_quota = MemoryQuota::new(capacity);
        let (mut tx, rx) = channel(buffer, memory_quota);
        let mut flag = true;
        let send = move |event| {
            flag = !flag;
            if flag {
                tx.unbounded_send(event, force_send)
            } else {
                block_on(tx.send_all(vec![event]))
            }
        };
        (Box::new(send), rx)
    }

    #[test]
    fn test_barrier() {
        let force_send = false;
        let (mut send, mut rx) = new_test_cancal(10, usize::MAX, force_send);
        send(CdcEvent::Event(Default::default())).unwrap();
        let (btx1, brx1) = mpsc::channel();
        send(CdcEvent::Barrier(Some(Box::new(move |()| {
            btx1.send(()).unwrap();
        }))))
        .unwrap();
        send(CdcEvent::ResolvedTs(Default::default())).unwrap();
        let (btx2, brx2) = mpsc::channel();
        send(CdcEvent::Barrier(Some(Box::new(move |()| {
            btx2.send(()).unwrap();
        }))))
        .unwrap();

        let mut drain = rx.drain();
        brx1.recv_timeout(Duration::from_millis(100)).unwrap_err();
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
        assert_matches!(block_on(drain.next()), Some((CdcEvent::Event(_), _)));
        brx1.recv_timeout(Duration::from_millis(100)).unwrap_err();
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
        assert_matches!(block_on(drain.next()), Some((CdcEvent::Barrier(_), _)));
        brx1.recv_timeout(Duration::from_millis(100)).unwrap();
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
        assert_matches!(block_on(drain.next()), Some((CdcEvent::ResolvedTs(_), _)));
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
        assert_matches!(block_on(drain.next()), Some((CdcEvent::Barrier(_), _)));
        brx2.recv_timeout(Duration::from_millis(100)).unwrap();
        brx1.recv_timeout(Duration::from_millis(100)).unwrap_err();
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
    }

    #[test]
    fn test_nonblocking_batch() {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let force_send = false;
        let (mut send, mut drain) =
            new_test_cancal(CDC_MSG_MAX_BATCH_SIZE * 2, usize::MAX, force_send);
        let (mut tx, mut rx) = unbounded();
        runtime.spawn(async move {
            drain.forward(&mut tx).await.unwrap();
        });
        for count in 1..CDC_EVENT_MAX_BATCH_SIZE + CDC_EVENT_MAX_BATCH_SIZE / 2 {
            for _ in 0..count {
                send(CdcEvent::Event(Default::default())).unwrap();
            }
            recv_timeout(&mut rx, Duration::from_millis(100)).unwrap();
        }

        if recv_timeout(&mut rx, Duration::from_millis(100)).is_ok() {
            panic!("expect to be timeout");
        }
    }

    #[test]
    fn test_congest() {
        let mut e = kvproto::cdcpb::Event::default();
        e.region_id = 1;
        let event = CdcEvent::Event(e.clone());
        assert!(event.size() != 0);
        // 1KB
        let max_pending_bytes = 1024;
        let buffer = max_pending_bytes / event.size();
        let force_send = false;
        let (mut send, _rx) = new_test_cancal(buffer as _, max_pending_bytes as _, force_send);
        for _ in 0..buffer {
            send(CdcEvent::Event(e.clone())).unwrap();
        }
        assert_matches!(send(CdcEvent::Event(e)).unwrap_err(), SendError::Congested);
    }

    #[test]
    fn test_force_send() {
        let mut e = kvproto::cdcpb::Event::default();
        e.region_id = 1;
        let event = CdcEvent::Event(e.clone());
        assert!(event.size() != 0);
        // 1KB
        let max_pending_bytes = 1024;
        let buffer = max_pending_bytes / event.size();
        let memory_quota = MemoryQuota::new(max_pending_bytes as _);
        let (tx, _rx) = channel(buffer as _, memory_quota);
        for _ in 0..buffer {
            tx.unbounded_send(CdcEvent::Event(e.clone()), false)
                .unwrap();
        }
        assert_matches!(
            tx.unbounded_send(CdcEvent::Event(e.clone()), false)
                .unwrap_err(),
            SendError::Congested
        );
        tx.unbounded_send(CdcEvent::Event(e), true).unwrap();
    }

    #[test]
    fn test_channel_memory_leak() {
        let mut e = kvproto::cdcpb::Event::default();
        e.region_id = 1;
        let event = CdcEvent::Event(e.clone());
        assert!(event.size() != 0);
        // 1KB
        let max_pending_bytes = 1024;
        let buffer = max_pending_bytes / event.size() + 1;
        let force_send = false;
        // Make sure memory quota is freed when rx is dropped before tx.
        {
            let (mut send, rx) = new_test_cancal(buffer as _, max_pending_bytes as _, force_send);
            loop {
                match send(CdcEvent::Event(e.clone())) {
                    Ok(_) => (),
                    Err(e) => {
                        assert_matches!(e, SendError::Congested);
                        break;
                    }
                }
            }
            let memory_quota = rx.memory_quota.clone();
            assert_eq!(memory_quota.alloc(event.size() as _), false,);
            drop(rx);
            assert_eq!(memory_quota.alloc(1024), true);
        }
        // Make sure memory quota is freed when tx is dropped before rx.
        {
            let (mut send, rx) = new_test_cancal(buffer as _, max_pending_bytes as _, force_send);
            loop {
                match send(CdcEvent::Event(e.clone())) {
                    Ok(_) => (),
                    Err(e) => {
                        assert_matches!(e, SendError::Congested);
                        break;
                    }
                }
            }
            let memory_quota = rx.memory_quota.clone();
            assert_eq!(memory_quota.alloc(event.size() as _), false,);
            drop(send);
            drop(rx);
            assert_eq!(memory_quota.alloc(1024), true);
        }
        // Make sure sending message to a closed channel does not leak memory quota.
        {
            let (mut send, rx) = new_test_cancal(buffer as _, max_pending_bytes as _, force_send);
            let memory_quota = rx.memory_quota.clone();
            assert_eq!(memory_quota.in_use(), 0);
            drop(rx);
            for _ in 0..max_pending_bytes {
                send(CdcEvent::Event(e.clone())).unwrap_err();
            }
            assert_eq!(memory_quota.in_use(), 0);
            assert_eq!(memory_quota.alloc(1024), true);
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use devtools_traits::{CONSOLE_API, CachedConsoleMessage, CachedConsoleMessageTypes, PAGE_ERROR};
use devtools_traits::{ComputedNodeLayout, ConsoleAPI, PageError, ScriptToDevtoolsControlMsg};
use devtools_traits::{EvaluateJSReply, Modification, NodeInfo, TimelineMarker, TimelineMarkerType};
use dom::bindings::codegen::Bindings::DOMRectBinding::{DOMRectMethods};
use dom::bindings::codegen::Bindings::DocumentBinding::DocumentMethods;
use dom::bindings::codegen::Bindings::ElementBinding::{ElementMethods};
use dom::bindings::conversions::{Castable, FromJSValConvertible, jsstring_to_str};
use dom::bindings::global::GlobalRef;
use dom::bindings::js::Root;
use dom::element::Element;
use dom::node::Node;
use ipc_channel::ipc::IpcSender;
use js::jsapi::{ObjectClassName, RootedObject, RootedValue};
use js::jsval::UndefinedValue;
use msg::constellation_msg::PipelineId;
use page::{IterablePage, Page};
use script_task::get_page;
use std::ffi::CStr;
use std::rc::Rc;
use std::str;
use uuid::Uuid;

#[allow(unsafe_code)]
pub fn handle_evaluate_js(global: &GlobalRef, eval: String, reply: IpcSender<EvaluateJSReply>) {
    let cx = global.get_cx();
    let mut rval = RootedValue::new(cx, UndefinedValue());
    global.evaluate_js_on_global_with_result(&eval, rval.handle_mut());

    reply.send(if rval.ptr.is_undefined() {
        EvaluateJSReply::VoidValue
    } else if rval.ptr.is_boolean() {
        EvaluateJSReply::BooleanValue(rval.ptr.to_boolean())
    } else if rval.ptr.is_double() || rval.ptr.is_int32() {
        EvaluateJSReply::NumberValue(
            FromJSValConvertible::from_jsval(cx, rval.handle(), ()).unwrap())
    } else if rval.ptr.is_string() {
        EvaluateJSReply::StringValue(jsstring_to_str(cx, rval.ptr.to_string()))
    } else if rval.ptr.is_null() {
        EvaluateJSReply::NullValue
    } else {
        assert!(rval.ptr.is_object());

        let obj = RootedObject::new(cx, rval.ptr.to_object());
        let class_name = unsafe { CStr::from_ptr(ObjectClassName(cx, obj.handle())) };
        let class_name = str::from_utf8(class_name.to_bytes()).unwrap();

        EvaluateJSReply::ActorValue {
            class: class_name.to_owned(),
            uuid: Uuid::new_v4().to_string(),
        }
    }).unwrap();
}

pub fn handle_get_root_node(page: &Rc<Page>, pipeline: PipelineId, reply: IpcSender<NodeInfo>) {
    let page = get_page(&*page, pipeline);
    let document = page.document();

    let node = document.upcast::<Node>();
    reply.send(node.summarize()).unwrap();
}

pub fn handle_get_document_element(page: &Rc<Page>, pipeline: PipelineId, reply: IpcSender<NodeInfo>) {
    let page = get_page(&*page, pipeline);
    let document = page.document();
    let document_element = document.r().GetDocumentElement().unwrap();

    let node = document_element.upcast::<Node>();
    reply.send(node.summarize()).unwrap();
}

fn find_node_by_unique_id(page: &Rc<Page>, pipeline: PipelineId, node_id: String) -> Root<Node> {
    let page = get_page(&*page, pipeline);
    let document = page.document();
    let node = document.upcast::<Node>();

    for candidate in node.traverse_preorder() {
        if candidate.r().get_unique_id() == node_id {
            return candidate;
        }
    }

    panic!("couldn't find node with unique id {}", node_id)
}

pub fn handle_get_children(page: &Rc<Page>, pipeline: PipelineId, node_id: String, reply: IpcSender<Vec<NodeInfo>>) {
    let parent = find_node_by_unique_id(&*page, pipeline, node_id);
    let children = parent.r().children().map(|child| {
        child.r().summarize()
    }).collect();
    reply.send(children).unwrap();
}

pub fn handle_get_layout(page: &Rc<Page>,
                         pipeline: PipelineId,
                         node_id: String,
                         reply: IpcSender<ComputedNodeLayout>) {
    let node = find_node_by_unique_id(&*page, pipeline, node_id);
    let elem = node.downcast::<Element>().expect("should be getting layout of element");
    let rect = elem.GetBoundingClientRect();
    let width = rect.Width() as f32;
    let height = rect.Height() as f32;
    reply.send(ComputedNodeLayout { width: width, height: height }).unwrap();
}

pub fn handle_get_cached_messages(_pipeline_id: PipelineId,
                                  message_types: CachedConsoleMessageTypes,
                                  reply: IpcSender<Vec<CachedConsoleMessage>>) {
    //TODO: check the messageTypes against a global Cache for console messages and page exceptions
    let mut messages = Vec::new();
    if message_types.contains(PAGE_ERROR) {
        //TODO: make script error reporter pass all reported errors
        //      to devtools and cache them for returning here.
        let msg = PageError {
            _type: "PageError".to_owned(),
            errorMessage: "page error test".to_owned(),
            sourceName: String::new(),
            lineText: String::new(),
            lineNumber: 0,
            columnNumber: 0,
            category: String::new(),
            timeStamp: 0,
            error: false,
            warning: false,
            exception: false,
            strict: false,
            private: false,
        };
        messages.push(CachedConsoleMessage::PageError(msg));
    }
    if message_types.contains(CONSOLE_API) {
        //TODO: do for real
        let msg = ConsoleAPI {
            _type: "ConsoleAPI".to_owned(),
            level: "error".to_owned(),
            filename: "http://localhost/~mihai/mozilla/test.html".to_owned(),
            lineNumber: 0,
            functionName: String::new(),
            timeStamp: 0,
            private: false,
            arguments: vec!["console error test".to_owned()],
        };
        messages.push(CachedConsoleMessage::ConsoleAPI(msg));
    }
    reply.send(messages).unwrap();
}

pub fn handle_modify_attribute(page: &Rc<Page>,
                               pipeline: PipelineId,
                               node_id: String,
                               modifications: Vec<Modification>) {
    let node = find_node_by_unique_id(&*page, pipeline, node_id);
    let elem = node.downcast::<Element>().expect("should be getting layout of element");

    for modification in &modifications {
        match modification.newValue {
            Some(ref string) => {
                let _ = elem.SetAttribute(modification.attributeName.clone(), string.clone());
            },
            None => elem.RemoveAttribute(modification.attributeName.clone()),
        }
    }
}

pub fn handle_wants_live_notifications(global: &GlobalRef, send_notifications: bool) {
    global.set_devtools_wants_updates(send_notifications);
}

pub fn handle_set_timeline_markers(page: &Rc<Page>,
                                   marker_types: Vec<TimelineMarkerType>,
                                   reply: IpcSender<TimelineMarker>) {
    let window = page.window();
    window.set_devtools_timeline_markers(marker_types, reply);
}

pub fn handle_drop_timeline_markers(page: &Rc<Page>,
                                    marker_types: Vec<TimelineMarkerType>) {
    let window = page.window();
    window.drop_devtools_timeline_markers(marker_types);
}

pub fn handle_request_animation_frame(page: &Rc<Page>, id: PipelineId, actor_name: String) {
    let page = page.find(id).expect("There is no such page");
    let doc = page.document();
    let devtools_sender = page.window().devtools_chan().unwrap();
    doc.r().request_animation_frame(box move |time| {
        let msg = ScriptToDevtoolsControlMsg::FramerateTick(actor_name, time);
        devtools_sender.send(msg).unwrap();
    });
}
<commit_msg>Avoid some string copies in handle_modify_attribute.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use devtools_traits::{CONSOLE_API, CachedConsoleMessage, CachedConsoleMessageTypes, PAGE_ERROR};
use devtools_traits::{ComputedNodeLayout, ConsoleAPI, PageError, ScriptToDevtoolsControlMsg};
use devtools_traits::{EvaluateJSReply, Modification, NodeInfo, TimelineMarker, TimelineMarkerType};
use dom::bindings::codegen::Bindings::DOMRectBinding::{DOMRectMethods};
use dom::bindings::codegen::Bindings::DocumentBinding::DocumentMethods;
use dom::bindings::codegen::Bindings::ElementBinding::{ElementMethods};
use dom::bindings::conversions::{Castable, FromJSValConvertible, jsstring_to_str};
use dom::bindings::global::GlobalRef;
use dom::bindings::js::Root;
use dom::element::Element;
use dom::node::Node;
use ipc_channel::ipc::IpcSender;
use js::jsapi::{ObjectClassName, RootedObject, RootedValue};
use js::jsval::UndefinedValue;
use msg::constellation_msg::PipelineId;
use page::{IterablePage, Page};
use script_task::get_page;
use std::ffi::CStr;
use std::rc::Rc;
use std::str;
use uuid::Uuid;

#[allow(unsafe_code)]
pub fn handle_evaluate_js(global: &GlobalRef, eval: String, reply: IpcSender<EvaluateJSReply>) {
    let cx = global.get_cx();
    let mut rval = RootedValue::new(cx, UndefinedValue());
    global.evaluate_js_on_global_with_result(&eval, rval.handle_mut());

    reply.send(if rval.ptr.is_undefined() {
        EvaluateJSReply::VoidValue
    } else if rval.ptr.is_boolean() {
        EvaluateJSReply::BooleanValue(rval.ptr.to_boolean())
    } else if rval.ptr.is_double() || rval.ptr.is_int32() {
        EvaluateJSReply::NumberValue(
            FromJSValConvertible::from_jsval(cx, rval.handle(), ()).unwrap())
    } else if rval.ptr.is_string() {
        EvaluateJSReply::StringValue(jsstring_to_str(cx, rval.ptr.to_string()))
    } else if rval.ptr.is_null() {
        EvaluateJSReply::NullValue
    } else {
        assert!(rval.ptr.is_object());

        let obj = RootedObject::new(cx, rval.ptr.to_object());
        let class_name = unsafe { CStr::from_ptr(ObjectClassName(cx, obj.handle())) };
        let class_name = str::from_utf8(class_name.to_bytes()).unwrap();

        EvaluateJSReply::ActorValue {
            class: class_name.to_owned(),
            uuid: Uuid::new_v4().to_string(),
        }
    }).unwrap();
}

pub fn handle_get_root_node(page: &Rc<Page>, pipeline: PipelineId, reply: IpcSender<NodeInfo>) {
    let page = get_page(&*page, pipeline);
    let document = page.document();

    let node = document.upcast::<Node>();
    reply.send(node.summarize()).unwrap();
}

pub fn handle_get_document_element(page: &Rc<Page>, pipeline: PipelineId, reply: IpcSender<NodeInfo>) {
    let page = get_page(&*page, pipeline);
    let document = page.document();
    let document_element = document.r().GetDocumentElement().unwrap();

    let node = document_element.upcast::<Node>();
    reply.send(node.summarize()).unwrap();
}

fn find_node_by_unique_id(page: &Rc<Page>, pipeline: PipelineId, node_id: String) -> Root<Node> {
    let page = get_page(&*page, pipeline);
    let document = page.document();
    let node = document.upcast::<Node>();

    for candidate in node.traverse_preorder() {
        if candidate.r().get_unique_id() == node_id {
            return candidate;
        }
    }

    panic!("couldn't find node with unique id {}", node_id)
}

pub fn handle_get_children(page: &Rc<Page>, pipeline: PipelineId, node_id: String, reply: IpcSender<Vec<NodeInfo>>) {
    let parent = find_node_by_unique_id(&*page, pipeline, node_id);
    let children = parent.r().children().map(|child| {
        child.r().summarize()
    }).collect();
    reply.send(children).unwrap();
}

pub fn handle_get_layout(page: &Rc<Page>,
                         pipeline: PipelineId,
                         node_id: String,
                         reply: IpcSender<ComputedNodeLayout>) {
    let node = find_node_by_unique_id(&*page, pipeline, node_id);
    let elem = node.downcast::<Element>().expect("should be getting layout of element");
    let rect = elem.GetBoundingClientRect();
    let width = rect.Width() as f32;
    let height = rect.Height() as f32;
    reply.send(ComputedNodeLayout { width: width, height: height }).unwrap();
}

pub fn handle_get_cached_messages(_pipeline_id: PipelineId,
                                  message_types: CachedConsoleMessageTypes,
                                  reply: IpcSender<Vec<CachedConsoleMessage>>) {
    //TODO: check the messageTypes against a global Cache for console messages and page exceptions
    let mut messages = Vec::new();
    if message_types.contains(PAGE_ERROR) {
        //TODO: make script error reporter pass all reported errors
        //      to devtools and cache them for returning here.
        let msg = PageError {
            _type: "PageError".to_owned(),
            errorMessage: "page error test".to_owned(),
            sourceName: String::new(),
            lineText: String::new(),
            lineNumber: 0,
            columnNumber: 0,
            category: String::new(),
            timeStamp: 0,
            error: false,
            warning: false,
            exception: false,
            strict: false,
            private: false,
        };
        messages.push(CachedConsoleMessage::PageError(msg));
    }
    if message_types.contains(CONSOLE_API) {
        //TODO: do for real
        let msg = ConsoleAPI {
            _type: "ConsoleAPI".to_owned(),
            level: "error".to_owned(),
            filename: "http://localhost/~mihai/mozilla/test.html".to_owned(),
            lineNumber: 0,
            functionName: String::new(),
            timeStamp: 0,
            private: false,
            arguments: vec!["console error test".to_owned()],
        };
        messages.push(CachedConsoleMessage::ConsoleAPI(msg));
    }
    reply.send(messages).unwrap();
}

pub fn handle_modify_attribute(page: &Rc<Page>,
                               pipeline: PipelineId,
                               node_id: String,
                               modifications: Vec<Modification>) {
    let node = find_node_by_unique_id(&*page, pipeline, node_id);
    let elem = node.downcast::<Element>().expect("should be getting layout of element");

    for modification in modifications {
        match modification.newValue {
            Some(string) => {
                let _ = elem.SetAttribute(modification.attributeName, string);
            },
            None => elem.RemoveAttribute(modification.attributeName),
        }
    }
}

pub fn handle_wants_live_notifications(global: &GlobalRef, send_notifications: bool) {
    global.set_devtools_wants_updates(send_notifications);
}

pub fn handle_set_timeline_markers(page: &Rc<Page>,
                                   marker_types: Vec<TimelineMarkerType>,
                                   reply: IpcSender<TimelineMarker>) {
    let window = page.window();
    window.set_devtools_timeline_markers(marker_types, reply);
}

pub fn handle_drop_timeline_markers(page: &Rc<Page>,
                                    marker_types: Vec<TimelineMarkerType>) {
    let window = page.window();
    window.drop_devtools_timeline_markers(marker_types);
}

pub fn handle_request_animation_frame(page: &Rc<Page>, id: PipelineId, actor_name: String) {
    let page = page.find(id).expect("There is no such page");
    let doc = page.document();
    let devtools_sender = page.window().devtools_chan().unwrap();
    doc.r().request_animation_frame(box move |time| {
        let msg = ScriptToDevtoolsControlMsg::FramerateTick(actor_name, time);
        devtools_sender.send(msg).unwrap();
    });
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::cell::DOMRefCell;
use dom::bindings::codegen::Bindings::BlobBinding;
use dom::bindings::codegen::Bindings::BlobBinding::BlobMethods;
use dom::bindings::codegen::UnionTypes::BlobOrString;
use dom::bindings::error::{Error, Fallible};
use dom::bindings::global::GlobalRef;
use dom::bindings::js::{JS, Root};
use dom::bindings::reflector::{Reflectable, Reflector, reflect_dom_object};
use dom::bindings::str::DOMString;
use encoding::all::UTF_8;
use encoding::types::{EncoderTrap, Encoding};
use ipc_channel::ipc;
use net_traits::{CoreResourceMsg, IpcSend};
use net_traits::blob_url_store::{BlobBuf, get_blob_origin};
use net_traits::filemanager_thread::{FileManagerThreadMsg, ReadFileProgress, RelativePos};
use std::cell::Cell;
use std::mem;
use std::ops::Index;
use std::path::PathBuf;
use uuid::Uuid;

/// File-based blob
#[derive(JSTraceable)]
pub struct FileBlob {
    id: Uuid,
    name: Option<PathBuf>,
    cache: DOMRefCell<Option<Vec<u8>>>,
    size: u64,
}

/// Different backends of Blob
#[must_root]
#[derive(JSTraceable)]
pub enum BlobImpl {
    /// File-based blob, whose content lives in the net process
    File(FileBlob),
    /// Memory-based blob, whose content lives in the script process
    Memory(Vec<u8>),
    /// Sliced blob, including parent blob reference and
    /// relative positions of current slicing range,
    /// IMPORTANT: The depth of tree is only two, i.e. the parent Blob must be
    /// either File-based or Memory-based
    Sliced(JS<Blob>, RelativePos),
}

impl BlobImpl {
    /// Construct memory-backed BlobImpl
    #[allow(unrooted_must_root)]
    pub fn new_from_bytes(bytes: Vec<u8>) -> BlobImpl {
        BlobImpl::Memory(bytes)
    }

    /// Construct file-backed BlobImpl from File ID
    pub fn new_from_file(file_id: Uuid, name: PathBuf, size: u64) -> BlobImpl {
        BlobImpl::File(FileBlob {
            id: file_id,
            name: Some(name),
            cache: DOMRefCell::new(None),
            size: size,
        })
    }
}

// https://w3c.github.io/FileAPI/#blob
#[dom_struct]
pub struct Blob {
    reflector_: Reflector,
    #[ignore_heap_size_of = "No clear owner"]
    blob_impl: DOMRefCell<BlobImpl>,
    /// content-type string
    type_string: String,
    isClosed_: Cell<bool>,
}

impl Blob {
    #[allow(unrooted_must_root)]
    pub fn new(global: GlobalRef, blob_impl: BlobImpl, typeString: String) -> Root<Blob> {
        let boxed_blob = box Blob::new_inherited(blob_impl, typeString);
        reflect_dom_object(boxed_blob, global, BlobBinding::Wrap)
    }

    #[allow(unrooted_must_root)]
    pub fn new_inherited(blob_impl: BlobImpl, type_string: String) -> Blob {
        Blob {
            reflector_: Reflector::new(),
            blob_impl: DOMRefCell::new(blob_impl),
            // NOTE: Guarding the format correctness here,
            // https://w3c.github.io/FileAPI/#dfn-type
            type_string: normalize_type_string(&type_string),
            isClosed_: Cell::new(false),
        }
    }

    #[allow(unrooted_must_root)]
    fn new_sliced(parent: &Blob, rel_pos: RelativePos,
                  relativeContentType: DOMString) -> Root<Blob> {
        let global = parent.global();
        let blob_impl = match *parent.blob_impl.borrow() {
            BlobImpl::File(_) => {
                // Create new parent node
                BlobImpl::Sliced(JS::from_ref(parent), rel_pos)
            }
            BlobImpl::Memory(_) => {
                // Create new parent node
                BlobImpl::Sliced(JS::from_ref(parent), rel_pos)
            }
            BlobImpl::Sliced(ref grandparent, ref old_rel_pos) => {
                // Adjust the slicing position, using same parent
                BlobImpl::Sliced(grandparent.clone(), old_rel_pos.slice_inner(&rel_pos))
            }
        };

        Blob::new(global.r(), blob_impl, relativeContentType.into())
    }

    // https://w3c.github.io/FileAPI/#constructorBlob
    pub fn Constructor(global: GlobalRef,
                       blobParts: Option<Vec<BlobOrString>>,
                       blobPropertyBag: &BlobBinding::BlobPropertyBag)
                       -> Fallible<Root<Blob>> {
        // TODO: accept other blobParts types - ArrayBuffer or ArrayBufferView
        let bytes: Vec<u8> = match blobParts {
            None => Vec::new(),
            Some(blobparts) => match blob_parts_to_bytes(blobparts) {
                Ok(bytes) => bytes,
                Err(_) => return Err(Error::InvalidCharacter),
            }
        };

        Ok(Blob::new(global, BlobImpl::new_from_bytes(bytes), blobPropertyBag.type_.to_string()))
    }

    /// Get a slice to inner data, this might incur synchronous read and caching
    pub fn get_bytes(&self) -> Result<Vec<u8>, ()> {
        match *self.blob_impl.borrow() {
            BlobImpl::File(ref f) => {
                let (buffer, is_new_buffer) = match *f.cache.borrow() {
                    Some(ref bytes) => (bytes.clone(), false),
                    None => {
                        let global = self.global();
                        let bytes = read_file(global.r(), f.id.clone())?;
                        (bytes, true)
                    }
                };

                // Cache
                if is_new_buffer {
                    *f.cache.borrow_mut() = Some(buffer.clone());
                }

                Ok(buffer)
            }
            BlobImpl::Memory(ref s) => Ok(s.clone()),
            BlobImpl::Sliced(ref parent, ref rel_pos) => {
                parent.get_bytes().map(|v| {
                    let range = rel_pos.to_abs_range(v.len());
                    v.index(range).to_vec()
                })
            }
        }
    }

    /// Get a FileID representing the Blob content,
    /// used by URL.createObjectURL
    pub fn get_blob_url_id(&self) -> Uuid {
        let opt_sliced_parent = match *self.blob_impl.borrow() {
            BlobImpl::Sliced(ref parent, ref rel_pos) => {
                Some((parent.promote(/* set_valid is */ false), rel_pos.clone(), parent.Size()))
            }
            _ => None
        };

        match opt_sliced_parent {
            Some((parent_id, rel_pos, size)) => self.create_sliced_url_id(&parent_id, &rel_pos, size),
            None => self.promote(/* set_valid is */ true),
        }
    }

    /// Promote non-Slice blob:
    /// 1. Memory-based: The bytes in data slice will be transferred to file manager thread.
    /// 2. File-based: If set_valid, then activate the FileID so it can serve as URL
    /// Depending on set_valid, the returned FileID can be part of
    /// valid or invalid Blob URL.
    fn promote(&self, set_valid: bool) -> Uuid {
        let mut bytes = vec![];

        match *self.blob_impl.borrow_mut() {
            BlobImpl::Sliced(_, _) => {
                debug!("Sliced can't have a sliced parent");
                // Return dummy id
                return Uuid::new_v4();
            }
            BlobImpl::File(ref f) => {
                if set_valid {
                    let global = self.global();
                    let origin = get_blob_origin(&global.r().get_url());
                    let (tx, rx) = ipc::channel().unwrap();

                    let msg = FileManagerThreadMsg::ActivateBlobURL(f.id.clone(), tx, origin.clone());
                    self.send_to_file_manager(msg);

                    match rx.recv().unwrap() {
                        Ok(_) => return f.id.clone(),
                        // Return a dummy id on error
                        Err(_) => return Uuid::new_v4(),
                    }
                } else {
                    // no need to activate
                    return f.id.clone();
                }
            }
            BlobImpl::Memory(ref mut bytes_in) => mem::swap(bytes_in, &mut bytes),
        };

        let global = self.global();
        let origin = get_blob_origin(&global.r().get_url());

        let blob_buf = BlobBuf {
            filename: None,
            type_string: self.type_string.clone(),
            size: bytes.len() as u64,
            bytes: bytes.to_vec(),
        };

        let (tx, rx) = ipc::channel().unwrap();
        let msg = FileManagerThreadMsg::PromoteMemory(blob_buf, set_valid, tx, origin.clone());
        self.send_to_file_manager(msg);

        match rx.recv().unwrap() {
            Ok(id) => {
                *self.blob_impl.borrow_mut() = BlobImpl::File(FileBlob {
                    id: id.clone(),
                    name: None,
                    cache: DOMRefCell::new(Some(bytes.to_vec())),
                    size: bytes.len() as u64,
                });
                id
            }
            // Dummy id
            Err(_) => Uuid::new_v4(),
        }
    }

    /// Get a FileID representing sliced parent-blob content
    fn create_sliced_url_id(&self, parent_id: &Uuid,
                            rel_pos: &RelativePos, parent_len: u64) -> Uuid {
        let global = self.global();

        let origin = get_blob_origin(&global.r().get_url());

        let (tx, rx) = ipc::channel().unwrap();
        let msg = FileManagerThreadMsg::AddSlicedURLEntry(parent_id.clone(),
                                                          rel_pos.clone(),
                                                          tx, origin.clone());
        self.send_to_file_manager(msg);
        match rx.recv().expect("File manager thread is down") {
            Ok(new_id) => {
                *self.blob_impl.borrow_mut() = BlobImpl::File(FileBlob {
                    id: new_id.clone(),
                    name: None,
                    cache: DOMRefCell::new(None),
                    size: rel_pos.to_abs_range(parent_len as usize).len() as u64,
                });

                // Return the indirect id reference
                new_id
            }
            Err(_) => {
                // Return dummy id
                Uuid::new_v4()
            }
        }
    }

    /// Cleanups at the time of destruction/closing
    fn clean_up_file_resource(&self) {
        if let BlobImpl::File(ref f) = *self.blob_impl.borrow() {
            let global = self.global();
            let origin = get_blob_origin(&global.r().get_url());

            let (tx, rx) = ipc::channel().unwrap();

            let msg = FileManagerThreadMsg::DecRef(f.id.clone(), origin, tx);
            self.send_to_file_manager(msg);
            let _ = rx.recv().unwrap();
        }
    }

    fn send_to_file_manager(&self, msg: FileManagerThreadMsg) {
        let global = self.global();
        let resource_threads = global.r().resource_threads();
        let _ = resource_threads.send(CoreResourceMsg::ToFileManager(msg));
    }
}

impl Drop for Blob {
    fn drop(&mut self) {
        if !self.IsClosed() {
            self.clean_up_file_resource();
        }
    }
}

fn read_file(global: GlobalRef, id: Uuid) -> Result<Vec<u8>, ()> {
    let resource_threads = global.resource_threads();
    let (chan, recv) = ipc::channel().map_err(|_|())?;
    let origin = get_blob_origin(&global.get_url());
    let check_url_validity = false;
    let msg = FileManagerThreadMsg::ReadFile(chan, id, check_url_validity, origin);
    let _ = resource_threads.send(CoreResourceMsg::ToFileManager(msg));

    let mut bytes = vec![];

    loop {
        match recv.recv().unwrap() {
            Ok(ReadFileProgress::Meta(mut blob_buf)) => {
                bytes.append(&mut blob_buf.bytes);
            }
            Ok(ReadFileProgress::Partial(mut bytes_in)) => {
                bytes.append(&mut bytes_in);
            }
            Ok(ReadFileProgress::EOF) => {
                return Ok(bytes);
            }
            Err(_) => return Err(()),
        }
    }
}

/// Extract bytes from BlobParts, used by Blob and File constructor
/// https://w3c.github.io/FileAPI/#constructorBlob
pub fn blob_parts_to_bytes(blobparts: Vec<BlobOrString>) -> Result<Vec<u8>, ()> {
    let mut ret = vec![];

    for blobpart in &blobparts {
        match blobpart {
            &BlobOrString::String(ref s) => {
                let mut bytes = UTF_8.encode(s, EncoderTrap::Replace).map_err(|_|())?;
                ret.append(&mut bytes);
            },
            &BlobOrString::Blob(ref b) => {
                let mut bytes = b.get_bytes().unwrap_or(vec![]);
                ret.append(&mut bytes);
            },
        }
    }

    Ok(ret)
}

impl BlobMethods for Blob {
    // https://w3c.github.io/FileAPI/#dfn-size
    fn Size(&self) -> u64 {
         match *self.blob_impl.borrow() {
            BlobImpl::File(ref f) => f.size,
            BlobImpl::Memory(ref v) => v.len() as u64,
            BlobImpl::Sliced(ref parent, ref rel_pos) =>
                rel_pos.to_abs_range(parent.Size() as usize).len() as u64,
         }
    }

    // https://w3c.github.io/FileAPI/#dfn-type
    fn Type(&self) -> DOMString {
        DOMString::from(self.type_string.clone())
    }

    // https://w3c.github.io/FileAPI/#slice-method-algo
    fn Slice(&self,
             start: Option<i64>,
             end: Option<i64>,
             contentType: Option<DOMString>)
             -> Root<Blob> {
        let rel_pos = RelativePos::from_opts(start, end);
        Blob::new_sliced(self, rel_pos, contentType.unwrap_or(DOMString::from("")))
    }

    // https://w3c.github.io/FileAPI/#dfn-isClosed
    fn IsClosed(&self) -> bool {
        self.isClosed_.get()
    }

    // https://w3c.github.io/FileAPI/#dfn-close
    fn Close(&self) {
        // Step 1
        if self.isClosed_.get() {
            return;
        }

        // Step 2
        self.isClosed_.set(true);

        // Step 3
        self.clean_up_file_resource();
    }
}

/// Get the normalized, MIME-parsable type string
/// https://w3c.github.io/FileAPI/#dfn-type
/// XXX: We will relax the restriction here,
/// since the spec has some problem over this part.
/// see https://github.com/w3c/FileAPI/issues/43
fn normalize_type_string(s: &str) -> String {
    if is_ascii_printable(s) {
        let s_lower = s.to_lowercase();
        // match s_lower.parse() as Result<Mime, ()> {
            // Ok(_) => s_lower,
            // Err(_) => "".to_string()
        s_lower
    } else {
        "".to_string()
    }
}

fn is_ascii_printable(string: &str) -> bool {
    // Step 5.1 in Sec 5.1 of File API spec
    // https://w3c.github.io/FileAPI/#constructorBlob
    string.chars().all(|c| c >= '\x20' && c <= '\x7E')
}
<commit_msg>Refactor isClosed_ to is_closed and relativeContentType to relative_content_type.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::cell::DOMRefCell;
use dom::bindings::codegen::Bindings::BlobBinding;
use dom::bindings::codegen::Bindings::BlobBinding::BlobMethods;
use dom::bindings::codegen::UnionTypes::BlobOrString;
use dom::bindings::error::{Error, Fallible};
use dom::bindings::global::GlobalRef;
use dom::bindings::js::{JS, Root};
use dom::bindings::reflector::{Reflectable, Reflector, reflect_dom_object};
use dom::bindings::str::DOMString;
use encoding::all::UTF_8;
use encoding::types::{EncoderTrap, Encoding};
use ipc_channel::ipc;
use net_traits::{CoreResourceMsg, IpcSend};
use net_traits::blob_url_store::{BlobBuf, get_blob_origin};
use net_traits::filemanager_thread::{FileManagerThreadMsg, ReadFileProgress, RelativePos};
use std::cell::Cell;
use std::mem;
use std::ops::Index;
use std::path::PathBuf;
use uuid::Uuid;

/// File-based blob
#[derive(JSTraceable)]
pub struct FileBlob {
    id: Uuid,
    name: Option<PathBuf>,
    cache: DOMRefCell<Option<Vec<u8>>>,
    size: u64,
}

/// Different backends of Blob
#[must_root]
#[derive(JSTraceable)]
pub enum BlobImpl {
    /// File-based blob, whose content lives in the net process
    File(FileBlob),
    /// Memory-based blob, whose content lives in the script process
    Memory(Vec<u8>),
    /// Sliced blob, including parent blob reference and
    /// relative positions of current slicing range,
    /// IMPORTANT: The depth of tree is only two, i.e. the parent Blob must be
    /// either File-based or Memory-based
    Sliced(JS<Blob>, RelativePos),
}

impl BlobImpl {
    /// Construct memory-backed BlobImpl
    #[allow(unrooted_must_root)]
    pub fn new_from_bytes(bytes: Vec<u8>) -> BlobImpl {
        BlobImpl::Memory(bytes)
    }

    /// Construct file-backed BlobImpl from File ID
    pub fn new_from_file(file_id: Uuid, name: PathBuf, size: u64) -> BlobImpl {
        BlobImpl::File(FileBlob {
            id: file_id,
            name: Some(name),
            cache: DOMRefCell::new(None),
            size: size,
        })
    }
}

// https://w3c.github.io/FileAPI/#blob
#[dom_struct]
pub struct Blob {
    reflector_: Reflector,
    #[ignore_heap_size_of = "No clear owner"]
    blob_impl: DOMRefCell<BlobImpl>,
    /// content-type string
    type_string: String,
    is_closed: Cell<bool>,
}

impl Blob {
    #[allow(unrooted_must_root)]
    pub fn new(global: GlobalRef, blob_impl: BlobImpl, typeString: String) -> Root<Blob> {
        let boxed_blob = box Blob::new_inherited(blob_impl, typeString);
        reflect_dom_object(boxed_blob, global, BlobBinding::Wrap)
    }

    #[allow(unrooted_must_root)]
    pub fn new_inherited(blob_impl: BlobImpl, type_string: String) -> Blob {
        Blob {
            reflector_: Reflector::new(),
            blob_impl: DOMRefCell::new(blob_impl),
            // NOTE: Guarding the format correctness here,
            // https://w3c.github.io/FileAPI/#dfn-type
            type_string: normalize_type_string(&type_string),
            is_closed: Cell::new(false),
        }
    }

    #[allow(unrooted_must_root)]
    fn new_sliced(parent: &Blob, rel_pos: RelativePos,
                  relative_content_type: DOMString) -> Root<Blob> {
        let global = parent.global();
        let blob_impl = match *parent.blob_impl.borrow() {
            BlobImpl::File(_) => {
                // Create new parent node
                BlobImpl::Sliced(JS::from_ref(parent), rel_pos)
            }
            BlobImpl::Memory(_) => {
                // Create new parent node
                BlobImpl::Sliced(JS::from_ref(parent), rel_pos)
            }
            BlobImpl::Sliced(ref grandparent, ref old_rel_pos) => {
                // Adjust the slicing position, using same parent
                BlobImpl::Sliced(grandparent.clone(), old_rel_pos.slice_inner(&rel_pos))
            }
        };

        Blob::new(global.r(), blob_impl, relative_content_type.into())
    }

    // https://w3c.github.io/FileAPI/#constructorBlob
    pub fn Constructor(global: GlobalRef,
                       blobParts: Option<Vec<BlobOrString>>,
                       blobPropertyBag: &BlobBinding::BlobPropertyBag)
                       -> Fallible<Root<Blob>> {
        // TODO: accept other blobParts types - ArrayBuffer or ArrayBufferView
        let bytes: Vec<u8> = match blobParts {
            None => Vec::new(),
            Some(blobparts) => match blob_parts_to_bytes(blobparts) {
                Ok(bytes) => bytes,
                Err(_) => return Err(Error::InvalidCharacter),
            }
        };

        Ok(Blob::new(global, BlobImpl::new_from_bytes(bytes), blobPropertyBag.type_.to_string()))
    }

    /// Get a slice to inner data, this might incur synchronous read and caching
    pub fn get_bytes(&self) -> Result<Vec<u8>, ()> {
        match *self.blob_impl.borrow() {
            BlobImpl::File(ref f) => {
                let (buffer, is_new_buffer) = match *f.cache.borrow() {
                    Some(ref bytes) => (bytes.clone(), false),
                    None => {
                        let global = self.global();
                        let bytes = read_file(global.r(), f.id.clone())?;
                        (bytes, true)
                    }
                };

                // Cache
                if is_new_buffer {
                    *f.cache.borrow_mut() = Some(buffer.clone());
                }

                Ok(buffer)
            }
            BlobImpl::Memory(ref s) => Ok(s.clone()),
            BlobImpl::Sliced(ref parent, ref rel_pos) => {
                parent.get_bytes().map(|v| {
                    let range = rel_pos.to_abs_range(v.len());
                    v.index(range).to_vec()
                })
            }
        }
    }

    /// Get a FileID representing the Blob content,
    /// used by URL.createObjectURL
    pub fn get_blob_url_id(&self) -> Uuid {
        let opt_sliced_parent = match *self.blob_impl.borrow() {
            BlobImpl::Sliced(ref parent, ref rel_pos) => {
                Some((parent.promote(/* set_valid is */ false), rel_pos.clone(), parent.Size()))
            }
            _ => None
        };

        match opt_sliced_parent {
            Some((parent_id, rel_pos, size)) => self.create_sliced_url_id(&parent_id, &rel_pos, size),
            None => self.promote(/* set_valid is */ true),
        }
    }

    /// Promote non-Slice blob:
    /// 1. Memory-based: The bytes in data slice will be transferred to file manager thread.
    /// 2. File-based: If set_valid, then activate the FileID so it can serve as URL
    /// Depending on set_valid, the returned FileID can be part of
    /// valid or invalid Blob URL.
    fn promote(&self, set_valid: bool) -> Uuid {
        let mut bytes = vec![];

        match *self.blob_impl.borrow_mut() {
            BlobImpl::Sliced(_, _) => {
                debug!("Sliced can't have a sliced parent");
                // Return dummy id
                return Uuid::new_v4();
            }
            BlobImpl::File(ref f) => {
                if set_valid {
                    let global = self.global();
                    let origin = get_blob_origin(&global.r().get_url());
                    let (tx, rx) = ipc::channel().unwrap();

                    let msg = FileManagerThreadMsg::ActivateBlobURL(f.id.clone(), tx, origin.clone());
                    self.send_to_file_manager(msg);

                    match rx.recv().unwrap() {
                        Ok(_) => return f.id.clone(),
                        // Return a dummy id on error
                        Err(_) => return Uuid::new_v4(),
                    }
                } else {
                    // no need to activate
                    return f.id.clone();
                }
            }
            BlobImpl::Memory(ref mut bytes_in) => mem::swap(bytes_in, &mut bytes),
        };

        let global = self.global();
        let origin = get_blob_origin(&global.r().get_url());

        let blob_buf = BlobBuf {
            filename: None,
            type_string: self.type_string.clone(),
            size: bytes.len() as u64,
            bytes: bytes.to_vec(),
        };

        let (tx, rx) = ipc::channel().unwrap();
        let msg = FileManagerThreadMsg::PromoteMemory(blob_buf, set_valid, tx, origin.clone());
        self.send_to_file_manager(msg);

        match rx.recv().unwrap() {
            Ok(id) => {
                *self.blob_impl.borrow_mut() = BlobImpl::File(FileBlob {
                    id: id.clone(),
                    name: None,
                    cache: DOMRefCell::new(Some(bytes.to_vec())),
                    size: bytes.len() as u64,
                });
                id
            }
            // Dummy id
            Err(_) => Uuid::new_v4(),
        }
    }

    /// Get a FileID representing sliced parent-blob content
    fn create_sliced_url_id(&self, parent_id: &Uuid,
                            rel_pos: &RelativePos, parent_len: u64) -> Uuid {
        let global = self.global();

        let origin = get_blob_origin(&global.r().get_url());

        let (tx, rx) = ipc::channel().unwrap();
        let msg = FileManagerThreadMsg::AddSlicedURLEntry(parent_id.clone(),
                                                          rel_pos.clone(),
                                                          tx, origin.clone());
        self.send_to_file_manager(msg);
        match rx.recv().expect("File manager thread is down") {
            Ok(new_id) => {
                *self.blob_impl.borrow_mut() = BlobImpl::File(FileBlob {
                    id: new_id.clone(),
                    name: None,
                    cache: DOMRefCell::new(None),
                    size: rel_pos.to_abs_range(parent_len as usize).len() as u64,
                });

                // Return the indirect id reference
                new_id
            }
            Err(_) => {
                // Return dummy id
                Uuid::new_v4()
            }
        }
    }

    /// Cleanups at the time of destruction/closing
    fn clean_up_file_resource(&self) {
        if let BlobImpl::File(ref f) = *self.blob_impl.borrow() {
            let global = self.global();
            let origin = get_blob_origin(&global.r().get_url());

            let (tx, rx) = ipc::channel().unwrap();

            let msg = FileManagerThreadMsg::DecRef(f.id.clone(), origin, tx);
            self.send_to_file_manager(msg);
            let _ = rx.recv().unwrap();
        }
    }

    fn send_to_file_manager(&self, msg: FileManagerThreadMsg) {
        let global = self.global();
        let resource_threads = global.r().resource_threads();
        let _ = resource_threads.send(CoreResourceMsg::ToFileManager(msg));
    }
}

impl Drop for Blob {
    fn drop(&mut self) {
        if !self.IsClosed() {
            self.clean_up_file_resource();
        }
    }
}

fn read_file(global: GlobalRef, id: Uuid) -> Result<Vec<u8>, ()> {
    let resource_threads = global.resource_threads();
    let (chan, recv) = ipc::channel().map_err(|_|())?;
    let origin = get_blob_origin(&global.get_url());
    let check_url_validity = false;
    let msg = FileManagerThreadMsg::ReadFile(chan, id, check_url_validity, origin);
    let _ = resource_threads.send(CoreResourceMsg::ToFileManager(msg));

    let mut bytes = vec![];

    loop {
        match recv.recv().unwrap() {
            Ok(ReadFileProgress::Meta(mut blob_buf)) => {
                bytes.append(&mut blob_buf.bytes);
            }
            Ok(ReadFileProgress::Partial(mut bytes_in)) => {
                bytes.append(&mut bytes_in);
            }
            Ok(ReadFileProgress::EOF) => {
                return Ok(bytes);
            }
            Err(_) => return Err(()),
        }
    }
}

/// Extract bytes from BlobParts, used by Blob and File constructor
/// https://w3c.github.io/FileAPI/#constructorBlob
pub fn blob_parts_to_bytes(blobparts: Vec<BlobOrString>) -> Result<Vec<u8>, ()> {
    let mut ret = vec![];

    for blobpart in &blobparts {
        match blobpart {
            &BlobOrString::String(ref s) => {
                let mut bytes = UTF_8.encode(s, EncoderTrap::Replace).map_err(|_|())?;
                ret.append(&mut bytes);
            },
            &BlobOrString::Blob(ref b) => {
                let mut bytes = b.get_bytes().unwrap_or(vec![]);
                ret.append(&mut bytes);
            },
        }
    }

    Ok(ret)
}

impl BlobMethods for Blob {
    // https://w3c.github.io/FileAPI/#dfn-size
    fn Size(&self) -> u64 {
         match *self.blob_impl.borrow() {
            BlobImpl::File(ref f) => f.size,
            BlobImpl::Memory(ref v) => v.len() as u64,
            BlobImpl::Sliced(ref parent, ref rel_pos) =>
                rel_pos.to_abs_range(parent.Size() as usize).len() as u64,
         }
    }

    // https://w3c.github.io/FileAPI/#dfn-type
    fn Type(&self) -> DOMString {
        DOMString::from(self.type_string.clone())
    }

    // https://w3c.github.io/FileAPI/#slice-method-algo
    fn Slice(&self,
             start: Option<i64>,
             end: Option<i64>,
             contentType: Option<DOMString>)
             -> Root<Blob> {
        let rel_pos = RelativePos::from_opts(start, end);
        Blob::new_sliced(self, rel_pos, contentType.unwrap_or(DOMString::from("")))
    }

    // https://w3c.github.io/FileAPI/#dfn-isClosed
    fn IsClosed(&self) -> bool {
        self.is_closed.get()
    }

    // https://w3c.github.io/FileAPI/#dfn-close
    fn Close(&self) {
        // Step 1
        if self.is_closed.get() {
            return;
        }

        // Step 2
        self.is_closed.set(true);

        // Step 3
        self.clean_up_file_resource();
    }
}

/// Get the normalized, MIME-parsable type string
/// https://w3c.github.io/FileAPI/#dfn-type
/// XXX: We will relax the restriction here,
/// since the spec has some problem over this part.
/// see https://github.com/w3c/FileAPI/issues/43
fn normalize_type_string(s: &str) -> String {
    if is_ascii_printable(s) {
        let s_lower = s.to_lowercase();
        // match s_lower.parse() as Result<Mime, ()> {
            // Ok(_) => s_lower,
            // Err(_) => "".to_string()
        s_lower
    } else {
        "".to_string()
    }
}

fn is_ascii_printable(string: &str) -> bool {
    // Step 5.1 in Sec 5.1 of File API spec
    // https://w3c.github.io/FileAPI/#constructorBlob
    string.chars().all(|c| c >= '\x20' && c <= '\x7E')
}
<|endoftext|>


<commit_before>
fn main () {
  // dummy
  println!("cargo:rustc-flags=-l sodium:static -L ../../libsodium-0.7.0/src/libsodium/.libs/");
}<commit_msg>Fixed libsodium-sys build script<commit_after>
fn main () {
  // dummy
  println!("cargo:rustc-flags=-l sodium");
}<|endoftext|>


<commit_before>use if_chain::if_chain;
use rustc_hir::{Local, PatKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_middle::lint::in_external_macro;
use rustc_middle::ty::subst::GenericArgKind;
use rustc_session::{declare_lint_pass, declare_tool_lint};

use crate::utils::{is_must_use_func_call, is_must_use_ty, match_type, paths, span_lint_and_help};

declare_clippy_lint! {
    /// **What it does:** Checks for `let _ = <expr>`
    /// where expr is #[must_use]
    ///
    /// **Why is this bad?** It's better to explicitly
    /// handle the value of a #[must_use] expr
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// ```rust
    /// fn f() -> Result<u32, u32> {
    ///     Ok(0)
    /// }
    ///
    /// let _ = f();
    /// // is_ok() is marked #[must_use]
    /// let _ = f().is_ok();
    /// ```
    pub LET_UNDERSCORE_MUST_USE,
    restriction,
    "non-binding let on a `#[must_use]` expression"
}

declare_clippy_lint! {
    /// **What it does:** Checks for `let _ = sync_lock`
    ///
    /// **Why is this bad?** This statement immediately drops the lock instead of
    /// extending it's lifetime to the end of the scope, which is often not intended.
    /// To extend lock lifetime to the end of the scope, use an underscore-prefixed
    /// name instead (i.e. _lock). If you want to explicitly drop the lock,
    /// `std::mem::drop` conveys your intention better and is less error-prone.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    ///
    /// Bad:
    /// ```rust,ignore
    /// let _ = mutex.lock();
    /// ```
    ///
    /// Good:
    /// ```rust,ignore
    /// let _lock = mutex.lock();
    /// ```
    pub LET_UNDERSCORE_LOCK,
    correctness,
    "non-binding let on a synchronization lock"
}

declare_lint_pass!(LetUnderscore => [LET_UNDERSCORE_MUST_USE, LET_UNDERSCORE_LOCK]);

const SYNC_GUARD_PATHS: [&[&str]; 3] = [
    &paths::MUTEX_GUARD,
    &paths::RWLOCK_READ_GUARD,
    &paths::RWLOCK_WRITE_GUARD,
];

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for LetUnderscore {
    fn check_local(&mut self, cx: &LateContext<'_, '_>, local: &Local<'_>) {
        if in_external_macro(cx.tcx.sess, local.span) {
            return;
        }

        if_chain! {
            if let PatKind::Wild = local.pat.kind;
            if let Some(ref init) = local.init;
            then {
                let init_ty = cx.tables.expr_ty(init);
                let contains_sync_guard = init_ty.walk().any(|inner| match inner.unpack() {
                    GenericArgKind::Type(inner_ty) => {
                        SYNC_GUARD_PATHS.iter().any(|path| match_type(cx, inner_ty, path))
                    },

                    GenericArgKind::Lifetime(_) | GenericArgKind::Const(_) => false,
                });
                if contains_sync_guard {
                    span_lint_and_help(
                        cx,
                        LET_UNDERSCORE_LOCK,
                        local.span,
                        "non-binding let on a synchronization lock",
                        None,
                        "consider using an underscore-prefixed named \
                            binding or dropping explicitly with `std::mem::drop`"
                    )
                } else if is_must_use_ty(cx, cx.tables.expr_ty(init)) {
                    span_lint_and_help(
                        cx,
                        LET_UNDERSCORE_MUST_USE,
                        local.span,
                        "non-binding let on an expression with `#[must_use]` type",
                        None,
                        "consider explicitly using expression value"
                    )
                } else if is_must_use_func_call(cx, init) {
                    span_lint_and_help(
                        cx,
                        LET_UNDERSCORE_MUST_USE,
                        local.span,
                        "non-binding let on a result of a `#[must_use]` function",
                        None,
                        "consider explicitly using function result"
                    )
                }
            }
        }
    }
}
<commit_msg>Rollup merge of #5726 - sozysozbot:patch-1, r=flip1995<commit_after>use if_chain::if_chain;
use rustc_hir::{Local, PatKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_middle::lint::in_external_macro;
use rustc_middle::ty::subst::GenericArgKind;
use rustc_session::{declare_lint_pass, declare_tool_lint};

use crate::utils::{is_must_use_func_call, is_must_use_ty, match_type, paths, span_lint_and_help};

declare_clippy_lint! {
    /// **What it does:** Checks for `let _ = <expr>`
    /// where expr is #[must_use]
    ///
    /// **Why is this bad?** It's better to explicitly
    /// handle the value of a #[must_use] expr
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// ```rust
    /// fn f() -> Result<u32, u32> {
    ///     Ok(0)
    /// }
    ///
    /// let _ = f();
    /// // is_ok() is marked #[must_use]
    /// let _ = f().is_ok();
    /// ```
    pub LET_UNDERSCORE_MUST_USE,
    restriction,
    "non-binding let on a `#[must_use]` expression"
}

declare_clippy_lint! {
    /// **What it does:** Checks for `let _ = sync_lock`
    ///
    /// **Why is this bad?** This statement immediately drops the lock instead of
    /// extending its lifetime to the end of the scope, which is often not intended.
    /// To extend lock lifetime to the end of the scope, use an underscore-prefixed
    /// name instead (i.e. _lock). If you want to explicitly drop the lock,
    /// `std::mem::drop` conveys your intention better and is less error-prone.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    ///
    /// Bad:
    /// ```rust,ignore
    /// let _ = mutex.lock();
    /// ```
    ///
    /// Good:
    /// ```rust,ignore
    /// let _lock = mutex.lock();
    /// ```
    pub LET_UNDERSCORE_LOCK,
    correctness,
    "non-binding let on a synchronization lock"
}

declare_lint_pass!(LetUnderscore => [LET_UNDERSCORE_MUST_USE, LET_UNDERSCORE_LOCK]);

const SYNC_GUARD_PATHS: [&[&str]; 3] = [
    &paths::MUTEX_GUARD,
    &paths::RWLOCK_READ_GUARD,
    &paths::RWLOCK_WRITE_GUARD,
];

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for LetUnderscore {
    fn check_local(&mut self, cx: &LateContext<'_, '_>, local: &Local<'_>) {
        if in_external_macro(cx.tcx.sess, local.span) {
            return;
        }

        if_chain! {
            if let PatKind::Wild = local.pat.kind;
            if let Some(ref init) = local.init;
            then {
                let init_ty = cx.tables.expr_ty(init);
                let contains_sync_guard = init_ty.walk().any(|inner| match inner.unpack() {
                    GenericArgKind::Type(inner_ty) => {
                        SYNC_GUARD_PATHS.iter().any(|path| match_type(cx, inner_ty, path))
                    },

                    GenericArgKind::Lifetime(_) | GenericArgKind::Const(_) => false,
                });
                if contains_sync_guard {
                    span_lint_and_help(
                        cx,
                        LET_UNDERSCORE_LOCK,
                        local.span,
                        "non-binding let on a synchronization lock",
                        None,
                        "consider using an underscore-prefixed named \
                            binding or dropping explicitly with `std::mem::drop`"
                    )
                } else if is_must_use_ty(cx, cx.tables.expr_ty(init)) {
                    span_lint_and_help(
                        cx,
                        LET_UNDERSCORE_MUST_USE,
                        local.span,
                        "non-binding let on an expression with `#[must_use]` type",
                        None,
                        "consider explicitly using expression value"
                    )
                } else if is_must_use_func_call(cx, init) {
                    span_lint_and_help(
                        cx,
                        LET_UNDERSCORE_MUST_USE,
                        local.span,
                        "non-binding let on a result of a `#[must_use]` function",
                        None,
                        "consider explicitly using function result"
                    )
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::{
    config::{ArtifactPersister, PersistConfig},
    errors::BuildProjectError,
    Artifact, ArtifactContent,
};
use lazy_static::lazy_static;
use log::debug;
use md5::{Digest, Md5};
use regex::Regex;
use std::{fs, path::PathBuf};

lazy_static! {
    static ref RELAY_HASH_REGEX: Regex = Regex::new(r#"@relayHash (\w{32})\n"#).unwrap();
    static ref REQUEST_ID_REGEX: Regex = Regex::new(r#"@relayRequestID (\w+)\n"#).unwrap();
}

pub async fn persist_operations(
    artifacts: &mut [Artifact],
    root_dir: &PathBuf,
    persist_config: &PersistConfig,
    artifact_persister: &Box<dyn ArtifactPersister + Send + Sync>,
) -> Result<(), BuildProjectError> {
    let handles = artifacts.iter_mut().map(|artifact| async move {
        if let ArtifactContent::Operation {
            ref text,
            ref mut id_and_text_hash,
            ..
        } = artifact.content
        {
            let text_hash = md5(text);
            let path = root_dir.join(&artifact.path);
            let persist_id = if let Some(extracted_id) = extract_persist_id(&path, &text_hash) {
                extracted_id
            } else {
                artifact_persister
                    .persist_artifact(text.to_string(), persist_config)
                    .await?
            };
            *id_and_text_hash = Some((persist_id, text_hash));
        }
        Ok(())
    });

    debug!("persisting {} documents", handles.len());
    let results = futures::future::join_all(handles).await;
    debug!("done persisting");
    let errors = results
        .into_iter()
        .filter_map(Result::err)
        .collect::<Vec<_>>();
    if !errors.is_empty() {
        return Err(BuildProjectError::PersistErrors { errors });
    }
    Ok(())
}

fn extract_persist_id(path: &PathBuf, text_hash: &str) -> Option<String> {
    let content = fs::read_to_string(path).ok()?;

    // Looks like a merge conflict, let's not trust this file.
    if content.contains("<<<<") || content.contains(">>>>") {
        return None;
    }

    // If the existing hash doesn't match the hash of the new query test, abort.
    if let Some(existing_hash) = extract_relay_hash(&content) {
        if existing_hash != text_hash {
            return None;
        }
    }

    extract_request_id(&content)
}

fn extract_relay_hash(content: &str) -> Option<&str> {
    RELAY_HASH_REGEX
        .captures(content)
        .and_then(|captures| captures.get(1).map(|m| m.as_str()))
}

fn extract_request_id(content: &str) -> Option<String> {
    REQUEST_ID_REGEX
        .captures(content)
        .map(|captures| captures[1].to_owned())
}

fn md5(data: &str) -> String {
    let mut md5 = Md5::new();
    md5.input(data);
    hex::encode(md5.result())
}
<commit_msg>only count operations we couldn't recover from disk<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::{
    config::{ArtifactPersister, PersistConfig},
    errors::BuildProjectError,
    Artifact, ArtifactContent,
};
use lazy_static::lazy_static;
use log::debug;
use md5::{Digest, Md5};
use regex::Regex;
use std::{fs, path::PathBuf};

lazy_static! {
    static ref RELAY_HASH_REGEX: Regex = Regex::new(r#"@relayHash (\w{32})\n"#).unwrap();
    static ref REQUEST_ID_REGEX: Regex = Regex::new(r#"@relayRequestID (\w+)\n"#).unwrap();
}

pub async fn persist_operations(
    artifacts: &mut [Artifact],
    root_dir: &PathBuf,
    persist_config: &PersistConfig,
    artifact_persister: &Box<dyn ArtifactPersister + Send + Sync>,
) -> Result<(), BuildProjectError> {
    let handles = artifacts
        .iter_mut()
        .flat_map(|artifact| {
            if let ArtifactContent::Operation {
                ref text,
                ref mut id_and_text_hash,
                ..
            } = artifact.content
            {
                let text_hash = md5(text);
                let artifact_path = root_dir.join(&artifact.path);
                if let Some(id) = extract_persist_id(&artifact_path, &text_hash) {
                    *id_and_text_hash = Some((id, text_hash));
                    None
                } else {
                    let text = text.clone();
                    Some(async move {
                        artifact_persister
                            .persist_artifact(text, persist_config)
                            .await
                            .map(|id| {
                                *id_and_text_hash = Some((id, text_hash));
                            })
                    })
                }
            } else {
                None
            }
        })
        .collect::<Vec<_>>();
    debug!("persisting {} documents", handles.len());
    let results = futures::future::join_all(handles).await;
    debug!("done persisting");
    let errors = results
        .into_iter()
        .filter_map(Result::err)
        .collect::<Vec<_>>();
    if !errors.is_empty() {
        return Err(BuildProjectError::PersistErrors { errors });
    }
    Ok(())
}

fn extract_persist_id(path: &PathBuf, text_hash: &str) -> Option<String> {
    let content = fs::read_to_string(path).ok()?;

    // Looks like a merge conflict, let's not trust this file.
    if content.contains("<<<<") || content.contains(">>>>") {
        return None;
    }

    // If the existing hash doesn't match the hash of the new query test, abort.
    if let Some(existing_hash) = extract_relay_hash(&content) {
        if existing_hash != text_hash {
            return None;
        }
    }

    extract_request_id(&content)
}

fn extract_relay_hash(content: &str) -> Option<&str> {
    RELAY_HASH_REGEX
        .captures(content)
        .and_then(|captures| captures.get(1).map(|m| m.as_str()))
}

fn extract_request_id(content: &str) -> Option<String> {
    REQUEST_ID_REGEX
        .captures(content)
        .map(|captures| captures[1].to_owned())
}

fn md5(data: &str) -> String {
    let mut md5 = Md5::new();
    md5.input(data);
    hex::encode(md5.result())
}
<|endoftext|>