use std::collections::VecDeque; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; use async_trait::async_trait; use log::trace; use futures::AsyncWriteExt; use futures::Stream; use kuska_handshake::async_std::BoxStreamWrite; use tokio::sync::mpsc; use crate::error::*; use crate::message::*; use crate::util::{ByteStream, Packet}; // Messages are sent by chunks // Chunk format: // - u32 BE: request id (same for request and response) // - u16 BE: chunk length, possibly with CHUNK_HAS_CONTINUATION flag // when this is not the last chunk of the message // - [u8; chunk_length] chunk data 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; struct SendQueueItem { id: RequestID, prio: RequestPriority, data: DataReader, } #[pin_project::pin_project] struct DataReader { #[pin] reader: ByteStream, packet: Packet, pos: usize, buf: Vec, eos: bool, } impl From for DataReader { fn from(data: ByteStream) -> DataReader { DataReader { reader: data, packet: Ok(Vec::new()), pos: 0, buf: Vec::with_capacity(MAX_CHUNK_LENGTH as usize), eos: false, } } } enum DataFrame { Data { /// a fixed size buffer containing some data, possibly padded with 0s data: [u8; MAX_CHUNK_LENGTH as usize], /// actual lenght of data len: usize, }, 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 { fn empty_last() -> Self { DataReaderItem { data: 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 len = match self.data { DataFrame::Data { len, .. } => len as u16, DataFrame::Error(e) => e as u16 | ERROR_MARKER, }; ChunkLength::to_be_bytes(len | continuation) } fn data(&self) -> &[u8] { match self.data { DataFrame::Data { ref data, len } => &data[..len], DataFrame::Error(_) => &[], } } } impl Stream for DataReader { type Item = DataReaderItem; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut this = self.project(); if *this.eos { // eos was reached at previous call to poll_next, where a partial packet // was returned. Now return None return Poll::Ready(None); } loop { let packet = match this.packet { Ok(v) => v, Err(e) => { let e = *e; *this.packet = Ok(Vec::new()); return Poll::Ready(Some(DataReaderItem { data: DataFrame::Error(e), may_have_more: true, })); } }; let packet_left = packet.len() - *this.pos; let buf_left = MAX_CHUNK_LENGTH as usize - this.buf.len(); let to_read = std::cmp::min(buf_left, packet_left); this.buf .extend_from_slice(&packet[*this.pos..*this.pos + to_read]); *this.pos += to_read; if this.buf.len() == MAX_CHUNK_LENGTH as usize { // we have a full buf, ready to send break; } // we don't have a full buf, packet is empty; try receive more if let Some(p) = futures::ready!(this.reader.as_mut().poll_next(cx)) { *this.packet = p; *this.pos = 0; // if buf is empty, we will loop and return the error directly. If buf // isn't empty, send it before by breaking. if this.packet.is_err() && !this.buf.is_empty() { break; } } else { *this.eos = true; break; } } let mut body = [0; MAX_CHUNK_LENGTH as usize]; 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 }, may_have_more: !*this.eos, })) } } struct SendQueue { items: VecDeque<(u8, VecDeque)>, } impl SendQueue { fn new() -> Self { Self { items: VecDeque::with_capacity(64), } } fn push(&mut self, item: SendQueueItem) { let prio = item.prio; let pos_prio = match self.items.binary_search_by(|(p, _)| p.cmp(&prio)) { Ok(i) => i, Err(i) => { self.items.insert(i, (prio, VecDeque::new())); i } }; self.items[pos_prio].1.push_back(item); } // used only in tests. They should probably be rewriten #[allow(dead_code)] fn pop(&mut self) -> Option { match self.items.pop_front() { None => None, Some((prio, mut items_at_prio)) => { let ret = items_at_prio.pop_front(); if !items_at_prio.is_empty() { self.items.push_front((prio, items_at_prio)); } ret.or_else(|| self.pop()) } } } fn is_empty(&self) -> bool { self.items.iter().all(|(_k, v)| v.is_empty()) } // this is like an async fn, but hand implemented fn next_ready(&mut self) -> SendQueuePollNextReady<'_> { SendQueuePollNextReady { queue: self } } } struct SendQueuePollNextReady<'a> { queue: &'a mut SendQueue, } impl<'a> futures::Future for SendQueuePollNextReady<'a> { type Output = (RequestID, DataReaderItem); fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll { 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(); 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())); } } } } // TODO what do we do if self.queue is empty? We won't get scheduled again. Poll::Pending } } /// The SendLoop trait, which is implemented both by the client and the server /// connection objects (ServerConna and ClientConn) adds a method `.send_loop()` /// that takes a channel of messages to send and an asynchronous writer, /// and sends messages from the channel to the async writer, putting them in a queue /// before being sent and doing the round-robin sending strategy. /// /// The `.send_loop()` exits when the sending end of the channel is closed, /// or if there is an error at any time writing to the async writer. #[async_trait] pub(crate) trait SendLoop: Sync { async fn send_loop( self: Arc, mut msg_recv: mpsc::UnboundedReceiver<(RequestID, RequestPriority, ByteStream)>, mut write: BoxStreamWrite, ) -> Result<(), Error> where W: AsyncWriteExt + Unpin + Send + Sync, { let mut sending = SendQueue::new(); let mut should_exit = false; while !should_exit || !sending.is_empty() { let recv_fut = msg_recv.recv(); futures::pin_mut!(recv_fut); let send_fut = sending.next_ready(); // recv_fut is cancellation-safe according to tokio doc, // send_fut is cancellation-safe as implemented above? use futures::future::Either; match futures::future::select(recv_fut, send_fut).await { Either::Left((sth, _send_fut)) => { if let Some((id, prio, data)) = sth { sending.push(SendQueueItem { id, prio, data: data.into(), }); } else { should_exit = true; }; } Either::Right(((id, data), _recv_fut)) => { trace!("send_loop: sending bytes for {}", id); let header_id = RequestID::to_be_bytes(id); write.write_all(&header_id[..]).await?; write.write_all(&data.header()).await?; write.write_all(data.data()).await?; write.flush().await?; } } } let _ = write.goodbye().await; Ok(()) } } #[cfg(test)] mod test { use super::*; fn empty_data() -> DataReader { type Item = Packet; let stream: Pin + Send + 'static>> = Box::pin(futures::stream::empty::()); stream.into() } #[test] fn test_priority_queue() { let i1 = SendQueueItem { id: 1, prio: PRIO_NORMAL, data: empty_data(), }; let i2 = SendQueueItem { id: 2, prio: PRIO_HIGH, data: empty_data(), }; let i2bis = SendQueueItem { id: 20, prio: PRIO_HIGH, data: empty_data(), }; let i3 = SendQueueItem { id: 3, prio: PRIO_HIGH | PRIO_SECONDARY, data: empty_data(), }; let i4 = SendQueueItem { id: 4, prio: PRIO_BACKGROUND | PRIO_SECONDARY, data: empty_data(), }; let i5 = SendQueueItem { id: 5, prio: PRIO_BACKGROUND | PRIO_PRIMARY, data: empty_data(), }; let mut q = SendQueue::new(); q.push(i1); // 1 let a = q.pop().unwrap(); // empty -> 1 assert_eq!(a.id, 1); assert!(q.pop().is_none()); q.push(a); // 1 q.push(i2); // 2 1 q.push(i2bis); // [2 20] 1 let a = q.pop().unwrap(); // 20 1 -> 2 assert_eq!(a.id, 2); let b = q.pop().unwrap(); // 1 -> 20 assert_eq!(b.id, 20); let c = q.pop().unwrap(); // empty -> 1 assert_eq!(c.id, 1); assert!(q.pop().is_none()); q.push(a); // 2 q.push(b); // [2 20] q.push(c); // [2 20] 1 q.push(i3); // [2 20] 3 1 q.push(i4); // [2 20] 3 1 4 q.push(i5); // [2 20] 3 1 5 4 let a = q.pop().unwrap(); // 20 3 1 5 4 -> 2 assert_eq!(a.id, 2); q.push(a); // [20 2] 3 1 5 4 let a = q.pop().unwrap(); // 2 3 1 5 4 -> 20 assert_eq!(a.id, 20); let b = q.pop().unwrap(); // 3 1 5 4 -> 2 assert_eq!(b.id, 2); q.push(b); // 2 3 1 5 4 let b = q.pop().unwrap(); // 3 1 5 4 -> 2 assert_eq!(b.id, 2); let c = q.pop().unwrap(); // 1 5 4 -> 3 assert_eq!(c.id, 3); q.push(b); // 2 1 5 4 let b = q.pop().unwrap(); // 1 5 4 -> 2 assert_eq!(b.id, 2); let e = q.pop().unwrap(); // 5 4 -> 1 assert_eq!(e.id, 1); let f = q.pop().unwrap(); // 4 -> 5 assert_eq!(f.id, 5); let g = q.pop().unwrap(); // empty -> 4 assert_eq!(g.id, 4); assert!(q.pop().is_none()); } }