use super::{ChannelError, MessageChannel};
use std::fmt;
use tokio::io::{
AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader,
};
use tokio::sync::Mutex;
const MAX_LINE: usize = 1 << 20;
pub struct CliMessageChannel {
io: Mutex<CliIo>,
}
impl fmt::Debug for CliMessageChannel {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("CliMessageChannel")
.field("io", &"Box<dyn AsyncBufRead> + Box<dyn AsyncWrite>")
.finish()
}
}
struct CliIo {
reader: Box<dyn AsyncBufRead + Unpin + Send>,
writer: Box<dyn AsyncWrite + Unpin + Send>,
}
impl CliMessageChannel {
pub fn new() -> Self {
Self::with_io(
Box::new(BufReader::new(tokio::io::stdin())),
Box::new(tokio::io::stdout()),
)
}
pub fn with_io(
reader: Box<dyn AsyncBufRead + Unpin + Send>,
writer: Box<dyn AsyncWrite + Unpin + Send>,
) -> Self {
Self {
io: Mutex::new(CliIo { reader, writer }),
}
}
}
impl Default for CliMessageChannel {
fn default() -> Self {
Self::new()
}
}
async fn write_line(
writer: &mut (dyn AsyncWrite + Unpin + Send),
message: &str,
) -> Result<(), ChannelError> {
writer
.write_all(message.as_bytes())
.await
.map_err(ChannelError::from)?;
writer.write_all(b"\n").await.map_err(ChannelError::from)?;
writer.flush().await.map_err(ChannelError::from)
}
#[async_trait::async_trait]
impl MessageChannel for CliMessageChannel {
async fn ask(&self, message: &str) -> Result<String, ChannelError> {
let mut io = self.io.lock().await;
write_line(io.writer.as_mut(), message).await?;
let mut line = String::new();
let read = io
.reader
.as_mut()
.take((MAX_LINE + 1) as u64)
.read_line(&mut line)
.await
.map_err(ChannelError::from)?;
if line.len() > MAX_LINE {
return Err(ChannelError::Io(format!(
"input line exceeds {MAX_LINE} bytes"
)));
}
if read == 0 {
return Err(ChannelError::Closed);
}
Ok(line.trim().to_string())
}
async fn notify(&self, message: &str) -> Result<(), ChannelError> {
let mut io = self.io.lock().await;
write_line(io.writer.as_mut(), message).await
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::pin::Pin;
use std::sync::{Arc, Mutex as StdMutex};
use std::task::{Context, Poll};
#[derive(Clone, Default)]
struct SharedBuf(Arc<StdMutex<Vec<u8>>>);
impl AsyncWrite for SharedBuf {
fn poll_write(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
self.0
.lock()
.expect("internal lock poisoned")
.extend_from_slice(buf);
Poll::Ready(Ok(buf.len()))
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Poll::Ready(Ok(()))
}
}
async fn channel(seed: &str) -> (CliMessageChannel, SharedBuf) {
let buf = SharedBuf(Arc::new(StdMutex::new(Vec::new())));
let (mut seed_tx, rx) = tokio::io::duplex(64);
seed_tx.write_all(seed.as_bytes()).await.unwrap();
let channel =
CliMessageChannel::with_io(Box::new(BufReader::new(rx)), Box::new(buf.clone()));
(channel, buf)
}
#[tokio::test]
async fn ask_reads_line_and_trims() {
let (channel, buf) = channel(" reply text \n").await;
assert_eq!(channel.ask("question").await.unwrap(), "reply text");
let out = String::from_utf8(buf.0.lock().unwrap().clone()).unwrap();
assert_eq!(out, "question\n");
}
#[tokio::test]
async fn ask_closed_input_returns_closed() {
let channel = CliMessageChannel::with_io(
Box::new(BufReader::new(tokio::io::empty())),
Box::new(tokio::io::sink()),
);
assert!(matches!(
channel.ask("question").await,
Err(ChannelError::Closed)
));
}
#[tokio::test]
async fn notify_writes_message() {
let (channel, buf) = channel("").await;
channel.notify("notice").await.unwrap();
let out = String::from_utf8(buf.0.lock().unwrap().clone()).unwrap();
assert_eq!(out, "notice\n");
}
#[tokio::test]
async fn concurrent_asks_serialized() {
let (channel, buf) = channel("r1\nr2\n").await;
let (a, b) = tokio::join!(channel.ask("m1"), channel.ask("m2"));
let mut replies = vec![a.unwrap(), b.unwrap()];
replies.sort();
assert_eq!(replies, vec!["r1", "r2"]);
let out = String::from_utf8(buf.0.lock().unwrap().clone()).unwrap();
let lines: Vec<&str> = out.lines().collect();
assert_eq!(lines.len(), 2);
assert!(lines.contains(&"m1") && lines.contains(&"m2"));
}
}