<commit_before>use byteorder::{BigEndian, ByteOrder};
use bytes::Bytes;
use futures::sync::{mpsc, BiLock};
use futures::{Async, Poll, Stream};
use std::collections::HashMap;
use std::time::Instant;

use util::SeqGenerator;

component! {
    ChannelManager : ChannelManagerInner {
        sequence: SeqGenerator<u16> = SeqGenerator::new(0),
        channels: HashMap<u16, mpsc::UnboundedSender<(u8, Bytes)>> = HashMap::new(),
        download_rate_estimate: usize = 0,
        download_measurement_start: Option<Instant> = None,
        download_measurement_bytes: usize = 0,
    }
}

#[derive(Debug, Hash, PartialEq, Eq, Copy, Clone)]
pub struct ChannelError;

pub struct Channel {
    receiver: mpsc::UnboundedReceiver<(u8, Bytes)>,
    state: ChannelState,
}

pub struct ChannelHeaders(BiLock<Channel>);
pub struct ChannelData(BiLock<Channel>);

pub enum ChannelEvent {
    Header(u8, Vec<u8>),
    Data(Bytes),
}

#[derive(Clone)]
enum ChannelState {
    Header(Bytes),
    Data,
    Closed,
}

impl ChannelManager {
    pub fn allocate(&self) -> (u16, Channel) {
        let (tx, rx) = mpsc::unbounded();

        let seq = self.lock(|inner| {
            let seq = inner.sequence.get();
            inner.channels.insert(seq, tx);
            seq
        });

        let channel = Channel {
            receiver: rx,
            state: ChannelState::Header(Bytes::new()),
        };

        (seq, channel)
    }

    pub(crate) fn dispatch(&self, cmd: u8, mut data: Bytes) {
        use std::collections::hash_map::Entry;

        let id: u16 = BigEndian::read_u16(data.split_to(2).as_ref());

        self.lock(|inner| {
            let current_time = Instant::now();
            if let Some(download_measurement_start) = inner.download_measurement_start {
                if (current_time - download_measurement_start).as_millis() > 1000 {
                    inner.download_rate_estimate = 1000 * inner.download_measurement_bytes
                        / (current_time - download_measurement_start).as_millis() as usize;
                    inner.download_measurement_start = Some(current_time);
                    inner.download_measurement_bytes = 0;
                }
            } else {
                inner.download_measurement_start = Some(current_time);
            }

            inner.download_measurement_bytes += data.len();

            if let Entry::Occupied(entry) = inner.channels.entry(id) {
                let _ = entry.get().unbounded_send((cmd, data));
            }
        });
    }

    pub fn get_download_rate_estimate(&self) -> usize {
        return self.lock(|inner| inner.download_rate_estimate);
    }
}

impl Channel {
    fn recv_packet(&mut self) -> Poll<Bytes, ChannelError> {
        let (cmd, packet) = match self.receiver.poll() {
            Ok(Async::Ready(t)) => t.expect("channel closed"),
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(()) => unreachable!(),
        };

        if cmd == 0xa {
            let code = BigEndian::read_u16(&packet.as_ref()[..2]);
            error!("channel error: {} {}", packet.len(), code);

            self.state = ChannelState::Closed;

            Err(ChannelError)
        } else {
            Ok(Async::Ready(packet))
        }
    }

    pub fn split(self) -> (ChannelHeaders, ChannelData) {
        let (headers, data) = BiLock::new(self);

        (ChannelHeaders(headers), ChannelData(data))
    }
}
<commit_msg>Fix error handling for closed channel.<commit_after>use byteorder::{BigEndian, ByteOrder};
use bytes::Bytes;
use futures::sync::{mpsc, BiLock};
use futures::{Async, Poll, Stream};
use std::collections::HashMap;
use std::time::Instant;

use util::SeqGenerator;

component! {
    ChannelManager : ChannelManagerInner {
        sequence: SeqGenerator<u16> = SeqGenerator::new(0),
        channels: HashMap<u16, mpsc::UnboundedSender<(u8, Bytes)>> = HashMap::new(),
        download_rate_estimate: usize = 0,
        download_measurement_start: Option<Instant> = None,
        download_measurement_bytes: usize = 0,
    }
}

#[derive(Debug, Hash, PartialEq, Eq, Copy, Clone)]
pub struct ChannelError;

pub struct Channel {
    receiver: mpsc::UnboundedReceiver<(u8, Bytes)>,
    state: ChannelState,
}

pub struct ChannelHeaders(BiLock<Channel>);
pub struct ChannelData(BiLock<Channel>);

pub enum ChannelEvent {
    Header(u8, Vec<u8>),
    Data(Bytes),
}

#[derive(Clone)]
enum ChannelState {
    Header(Bytes),
    Data,
    Closed,
}

impl ChannelManager {
    pub fn allocate(&self) -> (u16, Channel) {
        let (tx, rx) = mpsc::unbounded();

        let seq = self.lock(|inner| {
            let seq = inner.sequence.get();
            inner.channels.insert(seq, tx);
            seq
        });

        let channel = Channel {
            receiver: rx,
            state: ChannelState::Header(Bytes::new()),
        };

        (seq, channel)
    }

    pub(crate) fn dispatch(&self, cmd: u8, mut data: Bytes) {
        use std::collections::hash_map::Entry;

        let id: u16 = BigEndian::read_u16(data.split_to(2).as_ref());

        self.lock(|inner| {
            let current_time = Instant::now();
            if let Some(download_measurement_start) = inner.download_measurement_start {
                if (current_time - download_measurement_start).as_millis() > 1000 {
                    inner.download_rate_estimate = 1000 * inner.download_measurement_bytes
                        / (current_time - download_measurement_start).as_millis() as usize;
                    inner.download_measurement_start = Some(current_time);
                    inner.download_measurement_bytes = 0;
                }
            } else {
                inner.download_measurement_start = Some(current_time);
            }

            inner.download_measurement_bytes += data.len();

            if let Entry::Occupied(entry) = inner.channels.entry(id) {
                let _ = entry.get().unbounded_send((cmd, data));
            }
        });
    }

    pub fn get_download_rate_estimate(&self) -> usize {
        return self.lock(|inner| inner.download_rate_estimate);
    }
}

impl Channel {
    fn recv_packet(&mut self) -> Poll<Bytes, ChannelError> {
        let (cmd, packet) = match self.receiver.poll() {
            Ok(Async::Ready(Some(t))) => t,
            Ok(Async::Ready(None)) => return Err(ChannelError), // The channel has been closed.
            Ok(Async::NotReady) => return Ok(Async::NotReady),
            Err(()) => unreachable!(),
        };

        if cmd == 0xa {
            let code = BigEndian::read_u16(&packet.as_ref()[..2]);
            error!("channel error: {} {}", packet.len(), code);

            self.state = ChannelState::Closed;

            Err(ChannelError)
        } else {
            Ok(Async::Ready(packet))
        }
    }

    pub fn split(self) -> (ChannelHeaders, ChannelData) {
        let (headers, data) = BiLock::new(self);

        (ChannelHeaders(headers), ChannelData(data))
    }
}

{'fraction_non_alphanumeric': 0.1341533180778032, 'fraction_numerical': 0.009153318077803204, 'mean_word_length': 2.451628825271471, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 38, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 30, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2086470', 'n_tokens_mistral': 2153, 'n_tokens_neox': 2090, 'n_words': 420}