netapp/src/send.rs

311 lines
8.1 KiB
Rust
Raw Normal View History

2022-07-21 15:34:53 +00:00
use std::collections::VecDeque;
use std::pin::Pin;
2020-12-02 12:30:47 +00:00
use std::sync::Arc;
use std::task::{Context, Poll};
2020-12-02 12:30:47 +00:00
2022-07-21 15:34:53 +00:00
use async_trait::async_trait;
use bytes::Bytes;
use log::trace;
2020-12-02 12:30:47 +00:00
2022-07-21 15:34:53 +00:00
use futures::AsyncWriteExt;
2021-10-25 07:27:57 +00:00
use kuska_handshake::async_std::BoxStreamWrite;
use tokio::sync::mpsc;
2020-12-02 12:30:47 +00:00
use crate::error::*;
2022-07-21 15:34:53 +00:00
use crate::message::*;
2022-07-22 10:45:38 +00:00
use crate::stream::*;
2020-12-02 12:30:47 +00:00
// 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
2021-10-12 15:59:46 +00:00
pub(crate) type RequestID = u32;
2022-07-21 15:34:53 +00:00
pub(crate) type ChunkLength = u16;
2022-07-21 15:59:15 +00:00
2022-07-21 15:34:53 +00:00
pub(crate) const MAX_CHUNK_LENGTH: ChunkLength = 0x3FF0;
pub(crate) const ERROR_MARKER: ChunkLength = 0x4000;
pub(crate) const CHUNK_HAS_CONTINUATION: ChunkLength = 0x8000;
2020-12-02 12:30:47 +00:00
2022-07-22 11:23:42 +00:00
struct SendQueue {
items: VecDeque<(u8, VecDeque<SendQueueItem>)>,
}
2020-12-02 12:30:47 +00:00
struct SendQueueItem {
id: RequestID,
prio: RequestPriority,
2022-07-22 11:23:42 +00:00
data: ByteStreamReader,
2020-12-02 12:30:47 +00:00
}
impl SendQueue {
fn new() -> Self {
Self {
items: VecDeque::with_capacity(64),
2020-12-02 12:30:47 +00:00
}
}
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);
2020-12-02 12:30:47 +00:00
}
// used only in tests. They should probably be rewriten
#[allow(dead_code)]
2020-12-02 12:30:47 +00:00
fn pop(&mut self) -> Option<SendQueueItem> {
match self.items.pop_front() {
2020-12-02 12:30:47 +00:00
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));
2020-12-02 12:30:47 +00:00
}
2020-12-07 15:00:12 +00:00
ret.or_else(|| self.pop())
2020-12-02 12:30:47 +00:00
}
}
}
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> {
2022-07-21 15:59:15 +00:00
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];
2022-07-21 15:59:15 +00:00
let mut ready_item = None;
for (j, item) in items_at_prio.iter_mut().enumerate() {
2022-07-22 11:23:42 +00:00
let mut item_reader = item.data.read_exact_or_eos(MAX_CHUNK_LENGTH as usize);
match Pin::new(&mut item_reader).poll(ctx) {
2022-07-21 15:59:15 +00:00
Poll::Pending => (),
Poll::Ready(ready_v) => {
2022-07-22 11:23:42 +00:00
ready_item = Some((j, ready_v, item.data.eos()));
2022-07-21 15:59:15 +00:00
break;
}
}
}
2022-07-21 15:59:15 +00:00
2022-07-22 11:23:42 +00:00
if let Some((j, bytes_or_err, eos)) = ready_item {
let data_frame = match bytes_or_err {
Ok(bytes) => DataFrame::Data(bytes, !eos),
Err(e) => DataFrame::Error(match e {
ReadExactError::Stream(code) => code,
_ => unreachable!(),
}),
};
2022-07-21 15:59:15 +00:00
let item = items_at_prio.remove(j).unwrap();
let id = item.id;
2022-07-22 11:23:42 +00:00
if !eos {
2022-07-21 15:59:15 +00:00
items_at_prio.push_back(item);
} else if items_at_prio.is_empty() {
self.queue.items.remove(i);
}
2022-07-22 11:23:42 +00:00
return Poll::Ready((id, data_frame));
2022-07-21 15:59:15 +00:00
}
}
2022-07-22 11:23:42 +00:00
// 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
}
2020-12-02 12:30:47 +00:00
}
2022-07-22 11:23:42 +00:00
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(_) => &[],
}
}
}
2022-02-21 11:01:04 +00:00
/// 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.
2020-12-02 12:30:47 +00:00
#[async_trait]
pub(crate) trait SendLoop: Sync {
2020-12-07 17:07:55 +00:00
async fn send_loop<W>(
2020-12-02 12:30:47 +00:00
self: Arc<Self>,
2022-07-21 15:34:53 +00:00
mut msg_recv: mpsc::UnboundedReceiver<(RequestID, RequestPriority, ByteStream)>,
2021-10-25 07:27:57 +00:00
mut write: BoxStreamWrite<W>,
2020-12-07 17:07:55 +00:00
) -> Result<(), Error>
where
W: AsyncWriteExt + Unpin + Send + Sync,
2020-12-07 17:07:55 +00:00
{
2020-12-02 12:30:47 +00:00
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,
2022-07-22 11:23:42 +00:00
data: ByteStreamReader::new(data),
});
} else {
should_exit = true;
};
}
Either::Right(((id, data), _recv_fut)) => {
trace!("send_loop: sending bytes for {}", id);
2020-12-02 12:30:47 +00:00
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?;
}
2020-12-02 12:30:47 +00:00
}
}
2021-10-25 11:58:42 +00:00
let _ = write.goodbye().await;
2020-12-02 12:30:47 +00:00
Ok(())
}
}
#[cfg(test)]
mod test {
use super::*;
fn empty_data() -> DataReader {
2022-07-18 13:21:13 +00:00
type Item = Packet;
let stream: Pin<Box<dyn futures::Stream<Item = Item> + Send + 'static>> =
2022-07-18 13:21:13 +00:00
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());
}
}