Move things around

1 files changed, 410 insertions, 0 deletions

diff --git a/src/send.rs b/src/send.rs
new file mode 100644
index 0000000..330d41d
--- /dev/null
+++ b/src/send.rs
@@ -0,0 +1,410 @@
+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<u8>,
+	eos: bool,
+}
+
+impl From<ByteStream> 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<Option<Self::Item>> {
+		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<SendQueueItem>)>,
+}
+
+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<SendQueueItem> {
+		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<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();
+
+				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<W>(
+		self: Arc<Self>,
+		mut msg_recv: mpsc::UnboundedReceiver<(RequestID, RequestPriority, ByteStream)>,
+		mut write: BoxStreamWrite<W>,
+	) -> 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<Box<dyn futures::Stream<Item = Item> + Send + 'static>> =
+			Box::pin(futures::stream::empty::<Packet>());
+		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());
+	}
+}