netapp/src/send.rs

use std::collections::VecDeque;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use async_trait::async_trait;
use bytes::Bytes;
use log::*;

use futures::AsyncWriteExt;
use kuska_handshake::async_std::BoxStreamWrite;
use tokio::sync::mpsc;

use crate::error::*;
use crate::message::*;
use crate::stream::*;

// 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 SendQueue {
	items: Vec<(u8, VecDeque<SendQueueItem>)>,
}

struct SendQueueItem {
	id: RequestID,
	prio: RequestPriority,
	data: ByteStreamReader,
}

impl SendQueue {
	fn new() -> Self {
		Self {
			items: Vec::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);
	}
	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, DataFrame);

	fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
		for (i, (_prio, items_at_prio)) in self.queue.items.iter_mut().enumerate() {
			let mut ready_item = None;
			for (j, item) in items_at_prio.iter_mut().enumerate() {
				let mut item_reader = item.data.read_exact_or_eos(MAX_CHUNK_LENGTH as usize);
				match Pin::new(&mut item_reader).poll(ctx) {
					Poll::Pending => (),
					Poll::Ready(ready_v) => {
						ready_item = Some((j, ready_v));
						break;
					}
				}
			}

			if let Some((j, bytes_or_err)) = ready_item {
				let item = items_at_prio.remove(j).unwrap();
				let id = item.id;
				let eos = item.data.eos();

				let data_frame = match bytes_or_err {
					Ok(bytes) => {
						trace!(
							"send queue poll next ready: id {} eos {:?} bytes {}",
							id,
							eos,
							bytes.len()
						);
						DataFrame::Data(bytes, !eos)
					}
					Err(e) => DataFrame::Error(match e {
						ReadExactError::Stream(code) => {
							trace!(
								"send queue poll next ready: id {} eos {:?} ERROR {}",
								id,
								eos,
								code
							);
							code
						}
						_ => unreachable!(),
					}),
				};

				if !eos && !matches!(data_frame, DataFrame::Error(_)) {
					items_at_prio.push_back(item);
				} else if items_at_prio.is_empty() {
					self.queue.items.remove(i);
				}

				return Poll::Ready((id, data_frame));
			}
		}
		// If the queue is empty, this futures is eternally pending.
		// This is ok because we use it in a select with another future
		// that can interrupt it.
		Poll::Pending
	}
}

enum DataFrame {
	/// a fixed size buffer containing some data + a boolean indicating whether
	/// 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
	Data(Bytes, bool),
	/// An error code automatically signals the end of the stream
	Error(u8),
}

impl DataFrame {
	fn header(&self) -> [u8; 2] {
		let header_u16 = match self {
			DataFrame::Data(data, false) => data.len() as u16,
			DataFrame::Data(data, true) => data.len() as u16 | CHUNK_HAS_CONTINUATION,
			DataFrame::Error(e) => *e as u16 | ERROR_MARKER,
		};
		ChunkLength::to_be_bytes(header_u16)
	}

	fn data(&self) -> &[u8] {
		match self {
			DataFrame::Data(ref data, _) => &data[..],
			DataFrame::Error(_) => &[],
		}
	}
}

/// 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>,
		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 msg_recv = Some(msg_recv);
		while msg_recv.is_some() || !sending.is_empty() {
			debug!(
				"Sending: {:?}",
				sending
					.items
					.iter()
					.map(|(_, i)| i.iter().map(|x| x.id))
					.flatten()
					.collect::<Vec<_>>()
			);

			let recv_fut = async {
				if let Some(chan) = &mut msg_recv {
					chan.recv().await
				} else {
					futures::future::pending().await
				}
			};
			let send_fut = sending.next_ready();

			// recv_fut is cancellation-safe according to tokio doc,
			// send_fut is cancellation-safe as implemented above?
			tokio::select! {
				sth = recv_fut => {
					if let Some((id, prio, data)) = sth {
						trace!("send_loop: add stream {} to send", id);
						sending.push(SendQueueItem {
							id,
							prio,
							data: ByteStreamReader::new(data),
						});
					} else {
						msg_recv = None;
					};
				}
				(id, data) = send_fut => {
					trace!(
						"send_loop: id {}, send {} bytes, header_size {}",
						id,
						data.data().len(),
						hex::encode(data.header())
					);

					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(())
	}
}