use std::collections::HashMap;
use std::convert::TryInto;
use std::sync::Arc;
use bytes::Bytes;
use serde::Serialize;
use futures::sink::SinkExt;
use futures::stream::StreamExt;
use tokio::prelude::*;
use tokio::sync::{mpsc, oneshot, Barrier};
use tokio::task::JoinHandle;
use tokio_util::codec::{BytesCodec, FramedRead, FramedWrite};
use amp_serde::{ErrorResponse, OkResponse, Request};
use crate::frame::Response;
use crate::{Decoder, Error, Frame, RawFrame};
#[derive(Debug)]
pub struct DispatchRequest(pub Bytes, pub RawFrame, pub Option<ReplyTicket>);
struct ExpectReply {
tag: u64,
reply: oneshot::Sender<Response>,
barrier: Arc<tokio::sync::Barrier>,
}
type _FrameMaker = Box<dyn FnOnce(Option<Bytes>) -> Result<Vec<u8>, amp_serde::Error> + Send>;
struct FrameMaker(_FrameMaker);
impl std::fmt::Debug for FrameMaker {
fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
write!(fmt, "callback")
}
}
#[derive(Debug)]
enum WriteCmd {
Reply(Bytes),
Request(FrameMaker, Option<oneshot::Sender<Response>>),
}
#[derive(Debug)]
pub struct ReplyTicket {
tag: Option<Bytes>,
write_handle: mpsc::Sender<WriteCmd>,
}
impl ReplyTicket {
pub async fn ok<R: Serialize>(mut self, reply: R) -> Result<(), Error> {
let tag = self.tag.take().expect("Tag taken out of sequence");
let reply = amp_serde::to_bytes(OkResponse { tag, fields: reply })?;
self.write_handle
.send(WriteCmd::Reply(reply.into()))
.await?;
Ok(())
}
pub async fn error(
mut self,
code: Option<String>,
description: Option<String>,
) -> Result<(), Error> {
let tag = self.tag.take().expect("Tag taken out of sequence");
let reply = amp_serde::to_bytes(ErrorResponse {
tag,
code: code.unwrap_or_else(|| "UNKNOWN".into()),
description: description.unwrap_or_else(|| "".into()),
})?;
self.write_handle
.send(WriteCmd::Reply(reply.into()))
.await?;
Ok(())
}
}
impl Drop for ReplyTicket {
fn drop(&mut self) {
if let Some(tag) = self.tag.take() {
let mut write_handle = self.write_handle.clone();
let reply = amp_serde::to_bytes(ErrorResponse {
tag,
code: "UNKNOWN".into(),
description: "Request dropped without reply".into(),
})
.unwrap();
tokio::spawn(async move {
write_handle
.send(WriteCmd::Reply(reply.into()))
.await
.expect("error on drop")
});
}
}
}
#[derive(Clone)]
pub struct RequestSender(mpsc::Sender<WriteCmd>);
impl RequestSender {
pub async fn call_remote<Q: Serialize + Send + 'static>(
&mut self,
command: String,
request: Q,
) -> Result<RawFrame, Error> {
let (tx, rx) = oneshot::channel();
let frame = FrameMaker(Box::new(move |tag| {
amp_serde::to_bytes(Request {
tag,
command,
fields: request,
})
}));
self.0.send(WriteCmd::Request(frame, Some(tx))).await?;
rx.await?.map_err(|err| Error::Remote {
code: err.code,
description: err.description,
})
}
pub async fn call_remote_noreply<Q: Serialize + Send + 'static>(
&mut self,
command: String,
request: Q,
) -> Result<(), Error> {
let frame = FrameMaker(Box::new(move |tag| {
amp_serde::to_bytes(Request {
tag,
command,
fields: request,
})
}));
self.0.send(WriteCmd::Request(frame, None)).await?;
Ok(())
}
}
pub struct Handle {
write_res: JoinHandle<Result<(), Error>>,
read_res: JoinHandle<Result<(), Error>>,
write_loop_handle: Option<mpsc::Sender<WriteCmd>>,
shutdown: Option<oneshot::Sender<()>>,
}
impl Handle {
pub fn shutdown(&mut self) {
self.write_loop_handle = None;
if let Some(s) = self.shutdown.take() {
let _ = s.send(());
}
}
pub async fn join(mut self) -> Result<(), Error> {
self.write_loop_handle = None;
self.write_res.await.unwrap()?;
if let Some(s) = self.shutdown.take() {
let _ = s.send(());
}
self.read_res.await.unwrap()?;
Ok(())
}
pub fn request_sender(&self) -> Option<RequestSender> {
self.write_loop_handle.as_ref().cloned().map(RequestSender)
}
}
pub fn serve<R, W>(input: R, output: W) -> (Handle, mpsc::Receiver<DispatchRequest>)
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
let (write_tx, write_rx) = mpsc::channel::<WriteCmd>(32);
let (dispatch_tx, dispatch_rx) = mpsc::channel::<DispatchRequest>(32);
let (expect_tx, expect_rx) = mpsc::channel::<ExpectReply>(32);
let (shutdown_tx, shutdown_rx) = oneshot::channel();
let read_res = tokio::spawn(read_loop(
input,
shutdown_rx,
write_tx.clone(),
dispatch_tx,
expect_rx,
));
let write_res = tokio::spawn(write_loop(output, write_rx, expect_tx));
(
Handle {
write_res,
read_res,
write_loop_handle: Some(write_tx),
shutdown: Some(shutdown_tx),
},
dispatch_rx,
)
}
type ReplyMap = HashMap<u64, oneshot::Sender<Response>>;
async fn read_loop<R>(
input: R,
mut shutdown: oneshot::Receiver<()>,
mut write_tx: mpsc::Sender<WriteCmd>,
mut dispatch_tx: mpsc::Sender<DispatchRequest>,
mut expect_rx: mpsc::Receiver<ExpectReply>,
) -> Result<(), Error>
where
R: AsyncRead + Unpin,
{
let codec_in: Decoder<RawFrame> = Decoder::new();
let mut input = FramedRead::new(input, codec_in);
let mut reply_map = ReplyMap::new();
loop {
tokio::select! {
frame = input.next() => {
if let Some(frame) = frame {
dispatch_frame(frame?, &mut reply_map, &mut write_tx, &mut dispatch_tx).await?;
} else {
break;
}
}
expect = expect_rx.recv() => {
if let Some(expect) = expect {
reply_map.insert(expect.tag, expect.reply);
expect.barrier.wait().await;
}
}
_ = &mut shutdown => break
}
}
Ok(())
}
async fn dispatch_frame(
frame: RawFrame,
reply_map: &mut ReplyMap,
write_tx: &mut mpsc::Sender<WriteCmd>,
dispatch_tx: &mut mpsc::Sender<DispatchRequest>,
) -> Result<(), Error> {
match frame.try_into()? {
Frame::Request {
tag,
command,
fields,
} => {
let ticket = tag.map(|tag| ReplyTicket {
tag: Some(tag),
write_handle: write_tx.clone(),
});
let _ = dispatch_tx
.send(DispatchRequest(command, fields, ticket))
.await;
}
Frame::Response { tag, response } => {
let reply_tx = std::str::from_utf8(&tag)
.ok()
.and_then(|tag_str| u64::from_str_radix(tag_str, 16).ok())
.and_then(|tag_u64| reply_map.remove(&tag_u64))
.ok_or(Error::UnmatchedReply)?;
reply_tx.send(response).map_err(|_| Error::SendError)?;
}
}
Ok(())
}
async fn write_loop<W>(
output: W,
mut input: mpsc::Receiver<WriteCmd>,
mut expect_tx: mpsc::Sender<ExpectReply>,
) -> Result<(), Error>
where
W: AsyncWrite + Unpin,
{
let mut output = FramedWrite::new(output, BytesCodec::new());
let mut seqno: u64 = 0;
while let Some(msg) = input.next().await {
match msg {
WriteCmd::Reply(frame) => {
output.send(frame).await?;
}
WriteCmd::Request(request, reply) => {
if let Some(reply) = reply {
seqno += 1;
let barrier = Arc::new(Barrier::new(2));
let expect = ExpectReply {
tag: seqno,
reply,
barrier: barrier.clone(),
};
expect_tx.send(expect).await?;
barrier.wait().await;
output
.send(request.0(Some(format!("{:x}", seqno).into()))?.into())
.await?;
} else {
output.send(request.0(None)?.into()).await?;
}
}
}
}
Ok(())
}