Refactor send.rs
This commit is contained in:
parent
f35fa7d18d
commit
9dffa812c4
2 changed files with 61 additions and 60 deletions
|
@ -1,16 +1,13 @@
|
|||
use std::collections::HashMap;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use log::trace;
|
||||
|
||||
use futures::channel::mpsc::{unbounded, UnboundedReceiver, UnboundedSender};
|
||||
use futures::AsyncReadExt;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::error::*;
|
||||
|
||||
use crate::send::*;
|
||||
use crate::util::Packet;
|
||||
|
||||
|
|
116
src/send.rs
116
src/send.rs
|
@ -24,6 +24,7 @@ use crate::util::{ByteStream, Packet};
|
|||
|
||||
pub(crate) type RequestID = u32;
|
||||
pub(crate) type ChunkLength = u16;
|
||||
|
||||
pub(crate) const MAX_CHUNK_LENGTH: ChunkLength = 0x3FF0;
|
||||
pub(crate) const ERROR_MARKER: ChunkLength = 0x4000;
|
||||
pub(crate) const CHUNK_HAS_CONTINUATION: ChunkLength = 0x8000;
|
||||
|
@ -62,53 +63,58 @@ enum DataFrame {
|
|||
data: [u8; MAX_CHUNK_LENGTH as usize],
|
||||
/// actual lenght of data
|
||||
len: usize,
|
||||
/// whethere there may be more data comming from this stream. Can be used for some
|
||||
/// optimization. It's an error to set it to false if there is more data, but it is correct
|
||||
/// (albeit sub-optimal) to set it to true if there is nothing coming after
|
||||
may_have_more: bool,
|
||||
},
|
||||
/// An error code automatically signals the end of the stream
|
||||
Error(u8),
|
||||
}
|
||||
|
||||
struct DataReaderItem {
|
||||
data: DataFrame,
|
||||
/// whethere there may be more data comming from this stream. Can be used for some
|
||||
/// optimization. It's an error to set it to false if there is more data, but it is correct
|
||||
/// (albeit sub-optimal) to set it to true if there is nothing coming after
|
||||
may_have_more: bool,
|
||||
}
|
||||
|
||||
impl DataReaderItem {
|
||||
impl DataFrame {
|
||||
fn empty_last() -> Self {
|
||||
DataReaderItem {
|
||||
data: DataFrame::Data {
|
||||
data: [0; MAX_CHUNK_LENGTH as usize],
|
||||
len: 0,
|
||||
},
|
||||
DataFrame::Data {
|
||||
data: [0; MAX_CHUNK_LENGTH as usize],
|
||||
len: 0,
|
||||
may_have_more: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn header(&self) -> [u8; 2] {
|
||||
let continuation = if self.may_have_more {
|
||||
CHUNK_HAS_CONTINUATION
|
||||
} else {
|
||||
0
|
||||
let header_u16 = match self {
|
||||
DataFrame::Data {
|
||||
len,
|
||||
may_have_more: false,
|
||||
..
|
||||
} => *len as u16,
|
||||
DataFrame::Data {
|
||||
len,
|
||||
may_have_more: true,
|
||||
..
|
||||
} => *len as u16 | CHUNK_HAS_CONTINUATION,
|
||||
DataFrame::Error(e) => *e as u16 | ERROR_MARKER,
|
||||
};
|
||||
let len = match self.data {
|
||||
DataFrame::Data { len, .. } => len as u16,
|
||||
DataFrame::Error(e) => e as u16 | ERROR_MARKER,
|
||||
};
|
||||
|
||||
ChunkLength::to_be_bytes(len | continuation)
|
||||
ChunkLength::to_be_bytes(header_u16)
|
||||
}
|
||||
|
||||
fn data(&self) -> &[u8] {
|
||||
match self.data {
|
||||
DataFrame::Data { ref data, len } => &data[..len],
|
||||
match self {
|
||||
DataFrame::Data { ref data, len, .. } => &data[..*len],
|
||||
DataFrame::Error(_) => &[],
|
||||
}
|
||||
}
|
||||
|
||||
fn may_have_more(&self) -> bool {
|
||||
match self {
|
||||
DataFrame::Data { may_have_more, .. } => *may_have_more,
|
||||
DataFrame::Error(_) => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Stream for DataReader {
|
||||
type Item = DataReaderItem;
|
||||
type Item = DataFrame;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
let mut this = self.project();
|
||||
|
@ -125,10 +131,8 @@ impl Stream for DataReader {
|
|||
Err(e) => {
|
||||
let e = *e;
|
||||
*this.packet = Ok(Vec::new());
|
||||
return Poll::Ready(Some(DataReaderItem {
|
||||
data: DataFrame::Error(e),
|
||||
may_have_more: true,
|
||||
}));
|
||||
*this.eos = true;
|
||||
return Poll::Ready(Some(DataFrame::Error(e)));
|
||||
}
|
||||
};
|
||||
let packet_left = packet.len() - *this.pos;
|
||||
|
@ -161,8 +165,9 @@ impl Stream for DataReader {
|
|||
let len = this.buf.len();
|
||||
body[..len].copy_from_slice(this.buf);
|
||||
this.buf.clear();
|
||||
Poll::Ready(Some(DataReaderItem {
|
||||
data: DataFrame::Data { data: body, len },
|
||||
Poll::Ready(Some(DataFrame::Data {
|
||||
data: body,
|
||||
len,
|
||||
may_have_more: !*this.eos,
|
||||
}))
|
||||
}
|
||||
|
@ -218,38 +223,37 @@ struct SendQueuePollNextReady<'a> {
|
|||
}
|
||||
|
||||
impl<'a> futures::Future for SendQueuePollNextReady<'a> {
|
||||
type Output = (RequestID, DataReaderItem);
|
||||
type Output = (RequestID, DataFrame);
|
||||
|
||||
fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||
for i in 0..self.queue.items.len() {
|
||||
let (_prio, items_at_prio) = &mut self.queue.items[i];
|
||||
|
||||
for _ in 0..items_at_prio.len() {
|
||||
let mut item = items_at_prio.pop_front().unwrap();
|
||||
|
||||
let mut ready_item = None;
|
||||
for (j, item) in items_at_prio.iter_mut().enumerate() {
|
||||
match Pin::new(&mut item.data).poll_next(ctx) {
|
||||
Poll::Pending => items_at_prio.push_back(item),
|
||||
Poll::Ready(Some(data)) => {
|
||||
let id = item.id;
|
||||
if data.may_have_more {
|
||||
self.queue.push(item);
|
||||
} else {
|
||||
if items_at_prio.is_empty() {
|
||||
// this priority level is empty, remove it
|
||||
self.queue.items.remove(i);
|
||||
}
|
||||
}
|
||||
return Poll::Ready((id, data));
|
||||
}
|
||||
Poll::Ready(None) => {
|
||||
if items_at_prio.is_empty() {
|
||||
// this priority level is empty, remove it
|
||||
self.queue.items.remove(i);
|
||||
}
|
||||
return Poll::Ready((item.id, DataReaderItem::empty_last()));
|
||||
Poll::Pending => (),
|
||||
Poll::Ready(ready_v) => {
|
||||
ready_item = Some((j, ready_v));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some((j, ready_v)) = ready_item {
|
||||
let item = items_at_prio.remove(j).unwrap();
|
||||
let id = item.id;
|
||||
if ready_v
|
||||
.as_ref()
|
||||
.map(|data| data.may_have_more())
|
||||
.unwrap_or(false)
|
||||
{
|
||||
items_at_prio.push_back(item);
|
||||
} else if items_at_prio.is_empty() {
|
||||
self.queue.items.remove(i);
|
||||
}
|
||||
return Poll::Ready((id, ready_v.unwrap_or_else(DataFrame::empty_last)));
|
||||
}
|
||||
}
|
||||
// TODO what do we do if self.queue is empty? We won't get scheduled again.
|
||||
Poll::Pending
|
||||
|
|
Loading…
Reference in a new issue