librespot/core/src/channel.rs

202 lines
6.2 KiB
Rust
Raw Normal View History

use std::collections::HashMap;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Instant;
2017-01-19 22:45:24 +00:00
use byteorder::{BigEndian, ByteOrder};
use bytes::Bytes;
use futures_core::Stream;
use futures_util::lock::BiLock;
2021-04-17 09:55:57 +00:00
use futures_util::{ready, StreamExt};
use tokio::sync::mpsc;
2017-01-19 22:45:24 +00:00
use crate::util::SeqGenerator;
2017-01-19 22:45:24 +00:00
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,
invalid: bool = false,
2017-01-19 22:45:24 +00:00
}
}
2018-02-11 11:37:08 +00:00
#[derive(Debug, Hash, PartialEq, Eq, Copy, Clone)]
2017-01-19 22:45:24 +00:00
pub struct ChannelError;
pub struct Channel {
receiver: mpsc::UnboundedReceiver<(u8, Bytes)>,
2017-01-19 22:45:24 +00:00
state: ChannelState,
}
pub struct ChannelHeaders(BiLock<Channel>);
pub struct ChannelData(BiLock<Channel>);
pub enum ChannelEvent {
Header(u8, Vec<u8>),
Data(Bytes),
2017-01-19 22:45:24 +00:00
}
#[derive(Clone)]
enum ChannelState {
Header(Bytes),
2017-01-19 22:45:24 +00:00
Data,
Closed,
}
impl ChannelManager {
pub fn allocate(&self) -> (u16, Channel) {
let (tx, rx) = mpsc::unbounded_channel();
2017-01-19 22:45:24 +00:00
let seq = self.lock(|inner| {
let seq = inner.sequence.get();
if !inner.invalid {
inner.channels.insert(seq, tx);
}
2017-01-19 22:45:24 +00:00
seq
});
let channel = Channel {
receiver: rx,
state: ChannelState::Header(Bytes::new()),
2017-01-19 22:45:24 +00:00
};
(seq, channel)
}
2018-02-09 20:06:16 +00:00
pub(crate) fn dispatch(&self, cmd: u8, mut data: Bytes) {
2017-01-19 22:45:24 +00:00
use std::collections::hash_map::Entry;
let id: u16 = BigEndian::read_u16(data.split_to(2).as_ref());
2017-01-19 22:45:24 +00:00
self.lock(|inner| {
let current_time = Instant::now();
if let Some(download_measurement_start) = inner.download_measurement_start {
2019-11-07 21:58:17 +00:00
if (current_time - download_measurement_start).as_millis() > 1000 {
2019-11-11 07:22:41 +00:00
inner.download_rate_estimate = 1000 * inner.download_measurement_bytes
/ (current_time - download_measurement_start).as_millis() as usize;
2019-11-07 21:58:17 +00:00
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();
2017-01-19 22:45:24 +00:00
if let Entry::Occupied(entry) = inner.channels.entry(id) {
let _ = entry.get().send((cmd, data));
2017-01-19 22:45:24 +00:00
}
});
}
pub fn get_download_rate_estimate(&self) -> usize {
2021-03-01 02:37:22 +00:00
self.lock(|inner| inner.download_rate_estimate)
}
pub(crate) fn shutdown(&self) {
self.lock(|inner| {
inner.invalid = true;
// destroy the sending halves of the channels to signal everyone who is waiting for something.
inner.channels.clear();
});
}
2017-01-19 22:45:24 +00:00
}
impl Channel {
2021-01-21 20:49:39 +00:00
fn recv_packet(&mut self, cx: &mut Context<'_>) -> Poll<Result<Bytes, ChannelError>> {
2021-04-17 09:55:57 +00:00
let (cmd, packet) = ready!(self.receiver.poll_recv(cx)).ok_or(ChannelError)?;
2017-01-19 22:45:24 +00:00
if cmd == 0xa {
2017-01-29 17:54:32 +00:00
let code = BigEndian::read_u16(&packet.as_ref()[..2]);
2017-01-19 22:45:24 +00:00
error!("channel error: {} {}", packet.len(), code);
self.state = ChannelState::Closed;
2021-01-21 20:49:39 +00:00
Poll::Ready(Err(ChannelError))
2017-01-19 22:45:24 +00:00
} else {
2021-01-21 20:49:39 +00:00
Poll::Ready(Ok(packet))
2017-01-19 22:45:24 +00:00
}
}
pub fn split(self) -> (ChannelHeaders, ChannelData) {
let (headers, data) = BiLock::new(self);
(ChannelHeaders(headers), ChannelData(data))
}
}
impl Stream for Channel {
2021-01-21 20:49:39 +00:00
type Item = Result<ChannelEvent, ChannelError>;
2017-01-19 22:45:24 +00:00
2021-01-21 20:49:39 +00:00
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2017-01-19 22:45:24 +00:00
loop {
match self.state.clone() {
ChannelState::Closed => panic!("Polling already terminated channel"),
ChannelState::Header(mut data) => {
2021-03-01 02:37:22 +00:00
if data.is_empty() {
2021-04-17 09:55:57 +00:00
data = ready!(self.recv_packet(cx))?;
2017-01-19 22:45:24 +00:00
}
let length = BigEndian::read_u16(data.split_to(2).as_ref()) as usize;
2017-01-19 22:45:24 +00:00
if length == 0 {
assert_eq!(data.len(), 0);
self.state = ChannelState::Data;
} else {
let header_id = data.split_to(1).as_ref()[0];
let header_data = data.split_to(length - 1).as_ref().to_owned();
2017-01-19 22:45:24 +00:00
self.state = ChannelState::Header(data);
2017-01-29 17:54:32 +00:00
let event = ChannelEvent::Header(header_id, header_data);
2021-01-21 20:49:39 +00:00
return Poll::Ready(Some(Ok(event)));
2017-01-19 22:45:24 +00:00
}
}
ChannelState::Data => {
2021-04-17 09:55:57 +00:00
let data = ready!(self.recv_packet(cx))?;
2021-03-01 02:37:22 +00:00
if data.is_empty() {
2017-01-19 22:45:24 +00:00
self.receiver.close();
self.state = ChannelState::Closed;
2021-01-21 20:49:39 +00:00
return Poll::Ready(None);
2017-01-19 22:45:24 +00:00
} else {
2017-01-29 17:54:32 +00:00
let event = ChannelEvent::Data(data);
2021-01-21 20:49:39 +00:00
return Poll::Ready(Some(Ok(event)));
2017-01-19 22:45:24 +00:00
}
}
}
}
}
}
impl Stream for ChannelData {
2021-01-21 20:49:39 +00:00
type Item = Result<Bytes, ChannelError>;
2017-01-19 22:45:24 +00:00
2021-01-21 20:49:39 +00:00
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2021-04-17 09:55:57 +00:00
let mut channel = ready!(self.0.poll_lock(cx));
2017-01-19 22:45:24 +00:00
loop {
2021-04-17 09:55:57 +00:00
match ready!(channel.poll_next_unpin(cx)?) {
2017-01-19 22:45:24 +00:00
Some(ChannelEvent::Header(..)) => (),
2021-01-21 20:49:39 +00:00
Some(ChannelEvent::Data(data)) => return Poll::Ready(Some(Ok(data))),
None => return Poll::Ready(None),
2017-01-19 22:45:24 +00:00
}
}
}
}
impl Stream for ChannelHeaders {
2021-01-21 20:49:39 +00:00
type Item = Result<(u8, Vec<u8>), ChannelError>;
2017-01-19 22:45:24 +00:00
2021-01-21 20:49:39 +00:00
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2021-04-17 09:55:57 +00:00
let mut channel = ready!(self.0.poll_lock(cx));
2021-01-21 20:49:39 +00:00
2021-04-17 09:55:57 +00:00
match ready!(channel.poll_next_unpin(cx)?) {
2021-01-21 20:49:39 +00:00
Some(ChannelEvent::Header(id, data)) => Poll::Ready(Some(Ok((id, data)))),
2021-04-17 09:55:57 +00:00
_ => Poll::Ready(None),
2017-01-19 22:45:24 +00:00
}
}
}