orchid/orchid-base/src/reqnot.rs

use std::cell::RefCell;
use std::marker::PhantomData;
use std::pin::{Pin, pin};
use std::rc::Rc;
use std::{io, mem};

use async_fn_stream::try_stream;
use bound::Bound;
use derive_destructure::destructure;
use futures::channel::mpsc::{self, Receiver, Sender, channel};
use futures::channel::oneshot;
use futures::future::LocalBoxFuture;
use futures::lock::{Mutex, MutexGuard};
use futures::{
	AsyncRead, AsyncWrite, AsyncWriteExt, FutureExt, SinkExt, Stream, StreamExt, stream_select,
};
use hashbrown::HashMap;
use orchid_api_traits::{Decode, Encode, Request, UnderRoot};

use crate::localset::LocalSet;

#[must_use = "Receipts indicate that a required action has been performed within a function. \
	Most likely this should be returned somewhere."]
pub struct Receipt<'a>(PhantomData<&'a mut ()>);
impl Receipt<'_> {
	/// Only call this function from a custom implementation of [RepWriter]
	pub fn _new() -> Self { Self(PhantomData) }
}

/// Write guard to outbound for the purpose of serializing a request. Only one
/// can exist at a time. Dropping this object should panic.
pub trait ReqWriter<'a> {
	/// Access to the underlying channel. This may be buffered.
	fn writer(&mut self) -> Pin<&mut dyn AsyncWrite>;
	/// Finalize the request, release the outbound channel, then queue for the
	/// reply on the inbound channel.
	fn send(self: Box<Self>) -> LocalBoxFuture<'a, io::Result<Box<dyn RepReader<'a> + 'a>>>;
}

/// Write guard to inbound for the purpose of deserializing a reply. While held,
/// no inbound requests or other replies can be processed.
///
/// Dropping this object should panic even if [RepReader::finish] returns
/// synchronously, because the API isn't cancellation safe in general so it is a
/// programmer error in all cases to drop an object related to it without proper
/// cleanup.
pub trait RepReader<'a> {
	/// Access to the underlying channel. The length of the message is inferred
	/// from the number of bytes read so this must not be buffered.
	fn reader(&mut self) -> Pin<&mut dyn AsyncRead>;
	/// Finish reading the request
	fn finish(self: Box<Self>) -> LocalBoxFuture<'a, ()>;
}

/// Write guard to outbound for the purpose of serializing a notification.
///
/// Dropping this object should panic for the same reason [RepReader] panics
pub trait MsgWriter<'a> {
	/// Access to the underlying channel. This may be buffered.
	fn writer(&mut self) -> Pin<&mut dyn AsyncWrite>;
	/// Send the notification
	fn finish(self: Box<Self>) -> LocalBoxFuture<'a, io::Result<()>>;
}

/// For initiating outbound requests and notifications
pub trait Client {
	fn start_request(&self) -> LocalBoxFuture<'_, io::Result<Box<dyn ReqWriter<'_> + '_>>>;
	fn start_notif(&self) -> LocalBoxFuture<'_, io::Result<Box<dyn MsgWriter<'_> + '_>>>;
}

impl<T: Client + ?Sized> ClientExt for T {}
/// Extension trait with convenience methods that handle outbound request and
/// notif lifecycle and typing
#[allow(async_fn_in_trait)]
pub trait ClientExt: Client {
	async fn request<T: Request + UnderRoot<Root: Encode>>(&self, t: T) -> io::Result<T::Response> {
		let mut req = self.start_request().await?;
		t.into_root().encode(req.writer().as_mut()).await?;
		let mut rep = req.send().await?;
		let response = T::Response::decode(rep.reader()).await;
		rep.finish().await;
		response
	}
	async fn notify<T: UnderRoot<Root: Encode>>(&self, t: T) -> io::Result<()> {
		let mut notif = self.start_notif().await?;
		t.into_root().encode(notif.writer().as_mut()).await?;
		notif.finish().await?;
		Ok(())
	}
}

pub trait ReqReader<'a> {
	fn reader(&mut self) -> Pin<&mut dyn AsyncRead>;
	fn finish(self: Box<Self>) -> LocalBoxFuture<'a, Box<dyn ReqHandle<'a> + 'a>>;
}
impl<'a, T: ReqReader<'a> + ?Sized> ReqReaderExt<'a> for T {}
#[allow(async_fn_in_trait)]
pub trait ReqReaderExt<'a>: ReqReader<'a> {
	async fn read_req<R: Decode>(&mut self) -> io::Result<R> { R::decode(self.reader()).await }
	async fn reply<R: Request>(
		self: Box<Self>,
		req: impl Evidence<R>,
		rep: &R::Response,
	) -> io::Result<Receipt<'a>> {
		self.finish().await.reply(req, rep).await
	}
	async fn start_reply(self: Box<Self>) -> io::Result<Box<dyn RepWriter<'a> + 'a>> {
		self.finish().await.start_reply().await
	}
}

pub trait ReqHandle<'a> {
	fn start_reply(self: Box<Self>) -> LocalBoxFuture<'a, io::Result<Box<dyn RepWriter<'a> + 'a>>>;
}
impl<'a, T: ReqHandle<'a> + ?Sized> ReqHandleExt<'a> for T {}
#[allow(async_fn_in_trait)]
pub trait ReqHandleExt<'a>: ReqHandle<'a> {
	async fn reply<Req: Request>(
		self: Box<Self>,
		_: impl Evidence<Req>,
		rep: &Req::Response,
	) -> io::Result<Receipt<'a>> {
		let mut reply = self.start_reply().await?;
		rep.encode(reply.writer()).await?;
		reply.finish().await
	}
}

pub trait RepWriter<'a> {
	fn writer(&mut self) -> Pin<&mut dyn AsyncWrite>;
	fn finish(self: Box<Self>) -> LocalBoxFuture<'a, io::Result<Receipt<'a>>>;
}

pub trait MsgReader<'a> {
	fn reader(&mut self) -> Pin<&mut dyn AsyncRead>;
	fn finish(self: Box<Self>) -> LocalBoxFuture<'a, ()>;
}
impl<'a, T: ?Sized + MsgReader<'a>> MsgReaderExt<'a> for T {}
#[allow(async_fn_in_trait)]
pub trait MsgReaderExt<'a>: MsgReader<'a> {
	async fn read<N: Decode>(mut self: Box<Self>) -> io::Result<N> {
		let n = N::decode(self.reader()).await;
		self.finish().await;
		n
	}
}

/// A form of [Evidence] that doesn't require the value to be kept around
pub struct Witness<T>(PhantomData<T>);
impl<T> Witness<T> {
	pub fn of(_: &T) -> Self { Self(PhantomData) }
}
impl<T> Copy for Witness<T> {}
impl<T> Clone for Witness<T> {
	fn clone(&self) -> Self { *self }
}

/// A proxy for the type of a value either previously saved into a [Witness] or
/// still available.
pub trait Evidence<T> {}
impl<T> Evidence<T> for &'_ T {}
impl<T> Evidence<T> for Witness<T> {}

type IoRef<T> = Pin<Box<T>>;
type IoLock<T> = Rc<Mutex<Pin<Box<T>>>>;
type IoGuard<T> = Bound<MutexGuard<'static, Pin<Box<T>>>, IoLock<T>>;

/// An incoming request. This holds a lock on the ingress channel.
pub struct IoReqReader<'a> {
	prefix: &'a [u8],
	read: IoGuard<dyn AsyncRead>,
	write: &'a Mutex<IoRef<dyn AsyncWrite>>,
}
impl<'a> ReqReader<'a> for IoReqReader<'a> {
	fn reader(&mut self) -> Pin<&mut dyn AsyncRead> { self.read.as_mut() }
	fn finish(self: Box<Self>) -> LocalBoxFuture<'a, Box<dyn ReqHandle<'a> + 'a>> {
		Box::pin(async {
			Box::new(IoReqHandle { prefix: self.prefix, write: self.write }) as Box<dyn ReqHandle<'a>>
		})
	}
}
pub struct IoReqHandle<'a> {
	prefix: &'a [u8],
	write: &'a Mutex<IoRef<dyn AsyncWrite>>,
}
impl<'a> ReqHandle<'a> for IoReqHandle<'a> {
	fn start_reply(self: Box<Self>) -> LocalBoxFuture<'a, io::Result<Box<dyn RepWriter<'a> + 'a>>> {
		Box::pin(async move {
			let mut write = self.write.lock().await;
			write.as_mut().write_all(self.prefix).await?;
			Ok(Box::new(IoRepWriter { write }) as Box<dyn RepWriter<'a>>)
		})
	}
}
pub struct IoRepWriter<'a> {
	write: MutexGuard<'a, IoRef<dyn AsyncWrite>>,
}
impl<'a> RepWriter<'a> for IoRepWriter<'a> {
	fn writer(&mut self) -> Pin<&mut dyn AsyncWrite> { self.write.as_mut() }
	fn finish(mut self: Box<Self>) -> LocalBoxFuture<'a, io::Result<Receipt<'a>>> {
		Box::pin(async move {
			self.writer().flush().await?;
			Ok(Receipt(PhantomData))
		})
	}
}

pub struct IoMsgReader<'a> {
	_pd: PhantomData<&'a mut ()>,
	read: IoGuard<dyn AsyncRead>,
}
impl<'a> MsgReader<'a> for IoMsgReader<'a> {
	fn reader(&mut self) -> Pin<&mut dyn AsyncRead> { self.read.as_mut() }
	fn finish(self: Box<Self>) -> LocalBoxFuture<'static, ()> { Box::pin(async {}) }
}

#[derive(Debug)]
struct ReplySub {
	id: u64,
	ack: oneshot::Sender<()>,
	cb: oneshot::Sender<IoGuard<dyn AsyncRead>>,
}

struct IoClient {
	output: IoLock<dyn AsyncWrite>,
	id: Rc<RefCell<u64>>,
	subscribe: Rc<Sender<ReplySub>>,
}
impl IoClient {
	fn new(output: IoLock<dyn AsyncWrite>) -> (Receiver<ReplySub>, Self) {
		let (req, rep) = mpsc::channel(0);
		(rep, Self { output, id: Rc::new(RefCell::new(0)), subscribe: Rc::new(req) })
	}
	async fn lock_out(&self) -> IoGuard<dyn AsyncWrite> {
		Bound::async_new(self.output.clone(), async |o| o.lock().await).await
	}
}
impl Client for IoClient {
	fn start_notif(&self) -> LocalBoxFuture<'_, io::Result<Box<dyn MsgWriter<'_> + '_>>> {
		Box::pin(async {
			let mut o = self.lock_out().await;
			0u64.encode(o.as_mut()).await?;
			Ok(Box::new(IoNotifWriter { o }) as Box<dyn MsgWriter>)
		})
	}
	fn start_request(&self) -> LocalBoxFuture<'_, io::Result<Box<dyn ReqWriter<'_> + '_>>> {
		Box::pin(async {
			let id = {
				let mut id_g = self.id.borrow_mut();
				*id_g += 1;
				*id_g
			};
			let (cb, reply) = oneshot::channel();
			let (ack, got_ack) = oneshot::channel();
			self.subscribe.as_ref().clone().send(ReplySub { id, ack, cb }).await.unwrap();
			got_ack.await.unwrap();
			let mut w = self.lock_out().await;
			id.encode(w.as_mut()).await?;
			Ok(Box::new(IoReqWriter { reply, w }) as Box<dyn ReqWriter>)
		})
	}
}

struct IoReqWriter {
	reply: oneshot::Receiver<IoGuard<dyn AsyncRead>>,
	w: IoGuard<dyn AsyncWrite>,
}
impl<'a> ReqWriter<'a> for IoReqWriter {
	fn writer(&mut self) -> Pin<&mut dyn AsyncWrite> { self.w.as_mut() }
	fn send(self: Box<Self>) -> LocalBoxFuture<'a, io::Result<Box<dyn RepReader<'a> + 'a>>> {
		Box::pin(async {
			let Self { reply, mut w } = *self;
			w.flush().await?;
			mem::drop(w);
			let i = reply.await.expect("Client dropped before reply received");
			Ok(Box::new(IoRepReader { i }) as Box<dyn RepReader>)
		})
	}
}

struct IoRepReader {
	i: IoGuard<dyn AsyncRead>,
}
impl<'a> RepReader<'a> for IoRepReader {
	fn reader(&mut self) -> Pin<&mut dyn AsyncRead> { self.i.as_mut() }
	fn finish(self: Box<Self>) -> LocalBoxFuture<'static, ()> { Box::pin(async {}) }
}

#[derive(destructure)]
struct IoNotifWriter {
	o: IoGuard<dyn AsyncWrite>,
}
impl<'a> MsgWriter<'a> for IoNotifWriter {
	fn writer(&mut self) -> Pin<&mut dyn AsyncWrite> { self.o.as_mut() }
	fn finish(mut self: Box<Self>) -> LocalBoxFuture<'static, io::Result<()>> {
		Box::pin(async move { self.o.flush().await })
	}
}

pub struct CommCtx {
	exit: Sender<()>,
	o: Rc<Mutex<Pin<Box<dyn AsyncWrite>>>>,
}

impl CommCtx {
	pub async fn exit(self) -> io::Result<()> {
		self.o.lock().await.as_mut().close().await?;
		self.exit.clone().send(()).await.expect("quit channel dropped");
		Ok(())
	}
}

/// Establish bidirectional request-notification communication over a duplex
/// channel. The returned [IoClient] can be used for notifications immediately,
/// but requests can only be received while the future is running. The future
/// will only resolve when [CommCtx::quit] is called. The generic type
/// parameters are associated with the client and serve to ensure with a runtime
/// check that the correct message families are sent in the correct directions
/// across the channel.
pub fn io_comm(
	o: Pin<Box<dyn AsyncWrite>>,
	i: Pin<Box<dyn AsyncRead>>,
) -> (impl Client + 'static, CommCtx, IoCommServer) {
	let i = Rc::new(Mutex::new(i));
	let o = Rc::new(Mutex::new(o));
	let (onsub, client) = IoClient::new(o.clone());
	let (exit, onexit) = channel(1);
	(client, CommCtx { exit, o: o.clone() }, IoCommServer { o, i, onsub, onexit })
}
pub struct IoCommServer {
	o: Rc<Mutex<Pin<Box<dyn AsyncWrite>>>>,
	i: Rc<Mutex<Pin<Box<dyn AsyncRead>>>>,
	onsub: Receiver<ReplySub>,
	onexit: Receiver<()>,
}
impl IoCommServer {
	pub async fn listen(
		self,
		notif: impl for<'a> AsyncFn(Box<dyn MsgReader<'a> + 'a>) -> io::Result<()>,
		req: impl for<'a> AsyncFn(Box<dyn ReqReader<'a> + 'a>) -> io::Result<Receipt<'a>>,
	) -> io::Result<()> {
		let Self { o, i, onexit, onsub } = self;
		enum Event {
			Input(u64, IoGuard<dyn AsyncRead>),
			Sub(ReplySub),
			Exit,
		}
		let exiting = RefCell::new(false);
		let input_stream = try_stream(async |mut h| {
			loop {
				let mut g = Bound::async_new(i.clone(), async |i| i.lock().await).await;
				match u64::decode(g.as_mut()).await {
					Ok(id) => h.emit(Event::Input(id, g)).await,
					Err(e) => match e.kind() {
						io::ErrorKind::BrokenPipe
						| io::ErrorKind::ConnectionAborted
						| io::ErrorKind::UnexpectedEof => h.emit(Event::Exit).await,
						_ => return Err(e),
					},
				}
			}
		});
		let (mut add_pending_req, fork_future) = LocalSet::new();
		let mut fork_stream = pin!(fork_future.fuse().into_stream());
		let mut pending_replies = HashMap::new();
		'body: {
			let mut shared = pin!(stream_select!(
				pin!(input_stream) as Pin<&mut dyn Stream<Item = io::Result<Event>>>,
				onsub.map(|sub| Ok(Event::Sub(sub))),
				fork_stream.as_mut().map(|res| {
					res.map(|()| panic!("this substream cannot exit while the loop is running"))
				}),
				onexit.map(|()| Ok(Event::Exit)),
			));
			while let Some(next) = shared.next().await {
				match next {
					Err(e) => break 'body Err(e),
					Ok(Event::Exit) => {
						*exiting.borrow_mut() = true;
						let mut out = o.lock().await;
						out.as_mut().flush().await?;
						out.as_mut().close().await?;
						break;
					},
					Ok(Event::Sub(ReplySub { id, ack, cb })) => {
						pending_replies.insert(id, cb);
						ack.send(()).unwrap();
					},
					Ok(Event::Input(0, read)) => {
						let notif = &notif;
						let notif_job =
							async move { notif(Box::new(IoMsgReader { _pd: PhantomData, read })).await };
						add_pending_req.send(Box::pin(notif_job)).await.unwrap();
					},
					// MSB == 0 is a request, !id where MSB == 1 is the corresponding response
					Ok(Event::Input(id, read)) if (id & (1 << (u64::BITS - 1))) == 0 => {
						let (o, req) = (o.clone(), &req);
						let req_job = async move {
							let mut prefix = Vec::new();
							(!id).encode_vec(&mut prefix);
							let _ = req(Box::new(IoReqReader { prefix: &pin!(prefix), read, write: &o })).await;
							Ok(())
						};
						add_pending_req.send(Box::pin(req_job)).await.unwrap();
					},
					Ok(Event::Input(id, read)) => {
						let cb = pending_replies.remove(&!id).expect("Reply to unrecognized request");
						cb.send(read).unwrap_or_else(|_| panic!("Failed to send reply"));
					},
				}
			}
			Ok(())
		}?;
		mem::drop(add_pending_req);
		while let Some(next) = fork_stream.next().await {
			next?
		}
		Ok(())
	}
}

#[cfg(test)]
mod test {
	use std::cell::RefCell;

	use futures::channel::mpsc;
	use futures::{SinkExt, StreamExt, join};
	use orchid_api_derive::{Coding, Hierarchy};
	use orchid_api_traits::Request;
	use test_executors::spin_on;
	use unsync_pipe::pipe;

	use crate::logging::test::TestLogger;
	use crate::logging::with_logger;
	use crate::reqnot::{ClientExt, MsgReaderExt, ReqReaderExt, io_comm};

	#[derive(Clone, Debug, PartialEq, Coding, Hierarchy)]
	#[extendable]
	struct TestNotif(u64);

	#[test]
	fn notification() {
		let logger = TestLogger::new(async |s| eprint!("{s}"));
		spin_on(with_logger(logger, async {
			let (in1, out2) = pipe(1024);
			let (in2, out1) = pipe(1024);
			let (received, mut on_receive) = mpsc::channel(2);
			let (_, recv_ctx, recv_srv) = io_comm(Box::pin(in2), Box::pin(out2));
			let (sender, ..) = io_comm(Box::pin(in1), Box::pin(out1));
			join!(
				async {
					recv_srv
						.listen(
							async |notif| {
								received.clone().send(notif.read::<TestNotif>().await?).await.unwrap();
								Ok(())
							},
							async |_| panic!("Should receive notif, not request"),
						)
						.await
						.unwrap()
				},
				async {
					sender.notify(TestNotif(3)).await.unwrap();
					assert_eq!(on_receive.next().await, Some(TestNotif(3)));
					sender.notify(TestNotif(4)).await.unwrap();
					assert_eq!(on_receive.next().await, Some(TestNotif(4)));
					recv_ctx.exit().await.unwrap();
				}
			);
		}))
	}

	#[derive(Clone, Debug, Coding, Hierarchy)]
	#[extendable]
	struct DummyRequest(u64);
	impl Request for DummyRequest {
		type Response = u64;
	}

	#[test]
	fn request() {
		let logger = TestLogger::new(async |s| eprint!("{s}"));
		spin_on(with_logger(logger, async {
			let (in1, out2) = pipe(1024);
			let (in2, out1) = pipe(1024);
			let (_, srv_ctx, srv) = io_comm(Box::pin(in2), Box::pin(out2));
			let (client, client_ctx, client_srv) = io_comm(Box::pin(in1), Box::pin(out1));
			join!(
				async {
					srv
						.listen(
							async |_| panic!("No notifs expected"),
							async |mut req| {
								let val = req.read_req::<DummyRequest>().await?;
								req.reply(&val, &(val.0 + 1)).await
							},
						)
						.await
						.unwrap()
				},
				async {
					client_srv
						.listen(
							async |_| panic!("Not expecting ingress notif"),
							async |_| panic!("Not expecting ingress req"),
						)
						.await
						.unwrap()
				},
				async {
					let response = client.request(DummyRequest(5)).await.unwrap();
					assert_eq!(response, 6);
					srv_ctx.exit().await.unwrap();
					client_ctx.exit().await.unwrap();
				}
			);
		}))
	}

	#[test]
	fn exit() {
		let logger = TestLogger::new(async |s| eprint!("{s}"));
		spin_on(with_logger(logger, async {
			let (input1, output1) = pipe(1024);
			let (input2, output2) = pipe(1024);
			let (reply_client, reply_context, reply_server) =
				io_comm(Box::pin(input1), Box::pin(output2));
			let (req_client, req_context, req_server) = io_comm(Box::pin(input2), Box::pin(output1));
			let reply_context = RefCell::new(Some(reply_context));
			let (exit, onexit) = futures::channel::oneshot::channel::<()>();
			join!(
				async move {
					reply_server
						.listen(
							async |hand| {
								let _notif = hand.read::<TestNotif>().await.unwrap();
								let context = reply_context.borrow_mut().take().unwrap();
								context.exit().await?;
								Ok(())
							},
							async |mut hand| {
								let req = hand.read_req::<DummyRequest>().await?;
								hand.reply(&req, &(req.0 + 1)).await
							},
						)
						.await
						.unwrap();
					exit.send(()).unwrap();
					let _client = reply_client;
				},
				async move {
					req_server
						.listen(
							async |_| panic!("Only the other server expected notifs"),
							async |_| panic!("Only the other server expected requests"),
						)
						.await
						.unwrap();
					let _ctx = req_context;
				},
				async move {
					req_client.request(DummyRequest(0)).await.unwrap();
					req_client.notify(TestNotif(0)).await.unwrap();
					onexit.await.unwrap();
				}
			)
		}));
	}
}