Implement request cancellation

4 files changed, 159 insertions, 21 deletions

diff --git a/src/client.rs b/src/client.rs
index 9726125..d82c91e 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -1,7 +1,9 @@
 use std::collections::HashMap;
 use std::net::SocketAddr;
+use std::pin::Pin;
 use std::sync::atomic::{self, AtomicU32};
 use std::sync::{Arc, Mutex};
+use std::task::Poll;
 
 use arc_swap::ArcSwapOption;
 use async_trait::async_trait;
@@ -9,6 +11,7 @@ use bytes::Bytes;
 use log::{debug, error, trace};
 
 use futures::io::AsyncReadExt;
+use futures::Stream;
 use kuska_handshake::async_std::{handshake_client, BoxStream};
 use tokio::net::TcpStream;
 use tokio::select;
@@ -35,7 +38,7 @@ pub(crate) struct ClientConn {
 	pub(crate) remote_addr: SocketAddr,
 	pub(crate) peer_id: NodeID,
 
-	query_send: ArcSwapOption<mpsc::UnboundedSender<SendStream>>,
+	query_send: ArcSwapOption<mpsc::UnboundedSender<SendItem>>,
 
 	next_query_number: AtomicU32,
 	inflight: Mutex<HashMap<RequestID, oneshot::Sender<ByteStream>>>,
@@ -193,7 +196,9 @@ impl ClientConn {
 		#[cfg(feature = "telemetry")]
 		span.set_attribute(KeyValue::new("len_query_msg", req_msg_len as i64));
 
-		query_send.send((id, prio, req_order, req_stream))?;
+		query_send.send(SendItem::Stream(id, prio, req_order, req_stream))?;
+
+		let canceller = CancelOnDrop::new(id, query_send.as_ref().clone());
 
 		cfg_if::cfg_if! {
 			if #[cfg(feature = "telemetry")] {
@@ -205,6 +210,8 @@ impl ClientConn {
 			}
 		}
 
+		let stream = Box::pin(canceller.for_stream(stream));
+
 		let resp_enc = RespEnc::decode(stream).await?;
 		debug!("client: got response to request {} (path {})", id, path);
 		Resp::from_enc(resp_enc)
@@ -223,6 +230,63 @@ impl RecvLoop for ClientConn {
 			if ch.send(stream).is_err() {
 				debug!("Could not send request response, probably because request was interrupted. Dropping response.");
 			}
+		} else {
+			debug!("Got unexpected response to request {}, dropping it", id);
+		}
+	}
+}
+
+// ----
+
+struct CancelOnDrop {
+	id: RequestID,
+	query_send: mpsc::UnboundedSender<SendItem>,
+}
+
+impl CancelOnDrop {
+	fn new(id: RequestID, query_send: mpsc::UnboundedSender<SendItem>) -> Self {
+		Self { id, query_send }
+	}
+	fn for_stream(self, stream: ByteStream) -> CancelOnDropStream {
+		CancelOnDropStream {
+			cancel: Some(self),
+			stream: stream,
+		}
+	}
+}
+
+impl Drop for CancelOnDrop {
+	fn drop(&mut self) {
+		trace!("cancelling request {}", self.id);
+		let _ = self.query_send.send(SendItem::Cancel(self.id));
+	}
+}
+
+#[pin_project::pin_project]
+struct CancelOnDropStream {
+	cancel: Option<CancelOnDrop>,
+	#[pin]
+	stream: ByteStream,
+}
+
+impl Stream for CancelOnDropStream {
+	type Item = Packet;
+
+	fn poll_next(
+		self: Pin<&mut Self>,
+		cx: &mut std::task::Context<'_>,
+	) -> Poll<Option<Self::Item>> {
+		let this = self.project();
+		let res = this.stream.poll_next(cx);
+		if matches!(res, Poll::Ready(None)) {
+			if let Some(c) = this.cancel.take() {
+				std::mem::forget(c)
+			}
 		}
+		res
+	}
+
+	fn size_hint(&self) -> (usize, Option<usize>) {
+		self.stream.size_hint()
 	}
 }
diff --git a/src/recv.rs b/src/recv.rs
index ac93c4b..8909190 100644
--- a/src/recv.rs
+++ b/src/recv.rs
@@ -53,6 +53,7 @@ impl Drop for Sender {
 #[async_trait]
 pub(crate) trait RecvLoop: Sync + 'static {
 	fn recv_handler(self: &Arc<Self>, id: RequestID, stream: ByteStream);
+	fn cancel_handler(self: &Arc<Self>, _id: RequestID) {}
 
 	async fn recv_loop<R>(self: Arc<Self>, mut read: R, debug_name: String) -> Result<(), Error>
 	where
@@ -78,6 +79,18 @@ pub(crate) trait RecvLoop: Sync + 'static {
 			read.read_exact(&mut header_size[..]).await?;
 			let size = ChunkLength::from_be_bytes(header_size);
 
+			if size == CANCEL_REQUEST {
+				if let Some(mut stream) = streams.remove(&id) {
+					let _ = stream.send(Err(std::io::Error::new(
+						std::io::ErrorKind::Other,
+						"netapp: cancel requested",
+					)));
+					stream.end();
+				}
+				self.cancel_handler(id);
+				continue;
+			}
+
 			let has_cont = (size & CHUNK_HAS_CONTINUATION) != 0;
 			let is_error = (size & ERROR_MARKER) != 0;
 			let size = (size & CHUNK_LENGTH_MASK) as usize;
@@ -88,7 +101,10 @@ pub(crate) trait RecvLoop: Sync + 'static {
 				let kind = u8_to_io_errorkind(next_slice[0]);
 				let msg =
 					std::str::from_utf8(&next_slice[1..]).unwrap_or("<invalid utf8 error message>");
-				debug!("recv_loop({}): got id {}, error {:?}: {}", debug_name, id, kind, msg);
+				debug!(
+					"recv_loop({}): got id {}, error {:?}: {}",
+					debug_name, id, kind, msg
+				);
 				Some(Err(std::io::Error::new(kind, msg.to_string())))
 			} else {
 				trace!(
diff --git a/src/send.rs b/src/send.rs
index d927d98..780bbcf 100644
--- a/src/send.rs
+++ b/src/send.rs
@@ -22,6 +22,7 @@ use crate::stream::*;
 //		CHUNK_HAS_CONTINUATION when this is not the last chunk of the stream
 //		ERROR_MARKER if this chunk denotes an error
 //		(these two flags are exclusive, an error denotes the end of the stream)
+//		**special value** 0xFFFF indicates a CANCEL message
 // - [u8; chunk_length], either
 //   - if not error: chunk data
 //   - if error:
@@ -35,8 +36,14 @@ pub(crate) const MAX_CHUNK_LENGTH: ChunkLength = 0x3FF0;
 pub(crate) const ERROR_MARKER: ChunkLength = 0x4000;
 pub(crate) const CHUNK_HAS_CONTINUATION: ChunkLength = 0x8000;
 pub(crate) const CHUNK_LENGTH_MASK: ChunkLength = 0x3FFF;
+pub(crate) const CANCEL_REQUEST: ChunkLength = 0xFFFF;
 
-pub(crate) type SendStream = (RequestID, RequestPriority, Option<OrderTag>, ByteStream);
+pub(crate) enum SendItem {
+	Stream(RequestID, RequestPriority, Option<OrderTag>, ByteStream),
+	Cancel(RequestID),
+}
+
+// ----
 
 struct SendQueue {
 	items: Vec<(u8, SendQueuePriority)>,
@@ -71,6 +78,11 @@ impl SendQueue {
 		};
 		self.items[pos_prio].1.push(item);
 	}
+	fn remove(&mut self, id: RequestID) {
+		for (_, prioq) in self.items.iter_mut() {
+			prioq.remove(id);
+		}
+	}
 	fn is_empty(&self) -> bool {
 		self.items.iter().all(|(_k, v)| v.is_empty())
 	}
@@ -96,6 +108,16 @@ impl SendQueuePriority {
 		}
 		self.items.push_back(item);
 	}
+	fn remove(&mut self, id: RequestID) {
+		if let Some(i) = self.items.iter().position(|x| x.id == id) {
+			let item = self.items.remove(i).unwrap();
+			if let Some(OrderTag(stream, order)) = item.order_tag {
+				let order_vec = self.order.get_mut(&stream).unwrap();
+				let j = order_vec.iter().position(|x| *x == order).unwrap();
+				order_vec.remove(j).unwrap();
+			}
+		}
+	}
 	fn is_empty(&self) -> bool {
 		self.items.is_empty()
 	}
@@ -229,7 +251,7 @@ impl DataFrame {
 pub(crate) trait SendLoop: Sync {
 	async fn send_loop<W>(
 		self: Arc<Self>,
-		msg_recv: mpsc::UnboundedReceiver<SendStream>,
+		msg_recv: mpsc::UnboundedReceiver<SendItem>,
 		mut write: BoxStreamWrite<W>,
 		debug_name: String,
 	) -> Result<(), Error>
@@ -264,16 +286,27 @@ pub(crate) trait SendLoop: Sync {
 			tokio::select! {
 				biased;	// always read incomming channel first if it has data
 				sth = recv_fut => {
-					if let Some((id, prio, order_tag, data)) = sth {
-						trace!("send_loop({}): add stream {} to send", debug_name, id);
-						sending.push(SendQueueItem {
-							id,
-							prio,
-							order_tag,
-							data: ByteStreamReader::new(data),
-						});
-					} else {
-						msg_recv = None;
+					match sth {
+						Some(SendItem::Stream(id, prio, order_tag, data)) => {
+							trace!("send_loop({}): add stream {} to send", debug_name, id);
+							sending.push(SendQueueItem {
+								id,
+								prio,
+								order_tag,
+								data: ByteStreamReader::new(data),
+							})
+						}
+						Some(SendItem::Cancel(id)) => {
+							trace!("send_loop({}): cancelling {}", debug_name, id);
+							sending.remove(id);
+							let header_id = RequestID::to_be_bytes(id);
+							write.write_all(&header_id[..]).await?;
+							write.write_all(&ChunkLength::to_be_bytes(CANCEL_REQUEST)).await?;
+							write.flush().await?;
+						}
+						None => {
+							msg_recv = None;
+						}
 					};
 				}
 				(id, data) = send_fut => {
diff --git a/src/server.rs b/src/server.rs
index 2c12d9d..f9eb121 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1,5 +1,6 @@
+use std::collections::HashMap;
 use std::net::SocketAddr;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 use arc_swap::ArcSwapOption;
 use async_trait::async_trait;
@@ -53,7 +54,8 @@ pub(crate) struct ServerConn {
 
 	netapp: Arc<NetApp>,
 
-	resp_send: ArcSwapOption<mpsc::UnboundedSender<SendStream>>,
+	resp_send: ArcSwapOption<mpsc::UnboundedSender<SendItem>>,
+	running_handlers: Mutex<HashMap<RequestID, tokio::task::JoinHandle<()>>>,
 }
 
 impl ServerConn {
@@ -99,6 +101,7 @@ impl ServerConn {
 			remote_addr,
 			peer_id,
 			resp_send: ArcSwapOption::new(Some(Arc::new(resp_send))),
+			running_handlers: Mutex::new(HashMap::new()),
 		});
 
 		netapp.connected_as_server(peer_id, conn.clone());
@@ -174,10 +177,15 @@ impl SendLoop for ServerConn {}
 #[async_trait]
 impl RecvLoop for ServerConn {
 	fn recv_handler(self: &Arc<Self>, id: RequestID, stream: ByteStream) {
-		let resp_send = self.resp_send.load_full().unwrap();
+		let resp_send = match self.resp_send.load_full() {
+			Some(c) => c,
+			None => return,
+		};
+
+		let mut rh = self.running_handlers.lock().unwrap();
 
 		let self2 = self.clone();
-		tokio::spawn(async move {
+		let jh = tokio::spawn(async move {
 			debug!("server: recv_handler got {}", id);
 
 			let (prio, resp_enc_result) = match ReqEnc::decode(stream).await {
@@ -189,9 +197,26 @@ impl RecvLoop for ServerConn {
 
 			let (resp_stream, resp_order) = RespEnc::encode(resp_enc_result);
 			resp_send
-				.send((id, prio, resp_order, resp_stream))
+				.send(SendItem::Stream(id, prio, resp_order, resp_stream))
 				.log_err("ServerConn recv_handler send resp bytes");
-			Ok::<_, Error>(())
+
+			self2.running_handlers.lock().unwrap().remove(&id);
 		});
+
+		rh.insert(id, jh);
+	}
+
+	fn cancel_handler(self: &Arc<Self>, id: RequestID) {
+		trace!("received cancel for request {}", id);
+
+		// If the handler is still running, abort it now
+		if let Some(jh) = self.running_handlers.lock().unwrap().remove(&id) {
+			jh.abort();
+		}
+
+		// Inform the response sender that we don't need to send the response
+		if let Some(resp_send) = self.resp_send.load_full() {
+			let _ = resp_send.send(SendItem::Cancel(id));
+		}
 	}
 }