10 changed files with 273 additions and 522 deletions
Unified View
Diff Options
-
+10 -7client/src/control/ws.rs
-
+34 -42client/src/main.rs
-
+0 -310proto/src/core/channel.rs
-
+14 -2proto/src/core/frame.rs
-
+3 -2proto/src/core/mod.rs
-
+65 -35proto/src/core/worker.rs
-
+112 -97proto/src/server/client.rs
-
+1 -1proto/src/server/mod.rs
-
+9 -7proto/src/server/testing.rs
-
+25 -19proto/tests/connect.rs
+ 10
- 7
client/src/control/ws.rs
View File
+ 34
- 42
client/src/main.rs
View File
+ 0
- 310
proto/src/core/channel.rs
View File
| @ -1,310 +0,0 @@ | |||||
| use std::fmt::Debug; | |||||
| use std::future::Future; | |||||
| use std::io; | |||||
| use futures::stream::{Stream, StreamExt}; | |||||
| use log::debug; | |||||
| use thiserror::Error; | |||||
| use tokio::net::{ | |||||
| tcp::{OwnedReadHalf, OwnedWriteHalf}, | |||||
| TcpStream, | |||||
| }; | |||||
| use crate::core::frame::{FrameReader, FrameWriter}; | |||||
| use crate::core::value::{ValueDecode, ValueEncode}; | |||||
| /// An error that arose while exchanging messages over a `Channel`. | |||||
| #[derive(Debug, Error)] | |||||
| pub enum ChannelError { | |||||
| #[error("unexpected end of file")] | |||||
| UnexpectedEof, | |||||
| #[error("i/o error: {0}")] | |||||
| IOError(#[from] io::Error), | |||||
| } | |||||
| impl ChannelError { | |||||
| #[cfg(test)] | |||||
| pub fn is_unexpected_eof(&self) -> bool { | |||||
| match self { | |||||
| ChannelError::UnexpectedEof => true, | |||||
| _ => false, | |||||
| } | |||||
| } | |||||
| } | |||||
| /// A wrapper around a frame reader. Logically a part of `Channel`. | |||||
| /// | |||||
| /// Exists to provide `Channel` functionality that requires borrowing the | |||||
| /// channel's reader only. This allows borrowing both the reader and the writer | |||||
| /// at the same time in `Channel` without resorting to static methods. | |||||
| #[derive(Debug)] | |||||
| struct ChannelReader<ReadFrame> { | |||||
| inner: FrameReader<ReadFrame, OwnedReadHalf>, | |||||
| } | |||||
| impl<ReadFrame> ChannelReader<ReadFrame> | |||||
| where | |||||
| ReadFrame: ValueDecode + Debug, | |||||
| { | |||||
| async fn read(&mut self) -> io::Result<Option<ReadFrame>> { | |||||
| self.inner.read().await.map(|frame| { | |||||
| debug!("Channel: received frame: {:?}", frame); | |||||
| frame | |||||
| }) | |||||
| } | |||||
| async fn read_strict(&mut self) -> Result<ReadFrame, ChannelError> { | |||||
| match self.read().await? { | |||||
| None => Err(ChannelError::UnexpectedEof), | |||||
| Some(frame) => Ok(frame), | |||||
| } | |||||
| } | |||||
| } | |||||
| /// An asynchronous bidirectional message channel over TCP. | |||||
| #[derive(Debug)] | |||||
| pub struct Channel<ReadFrame, WriteFrame> { | |||||
| reader: ChannelReader<ReadFrame>, | |||||
| writer: FrameWriter<WriteFrame, OwnedWriteHalf>, | |||||
| } | |||||
| impl<ReadFrame, WriteFrame> Channel<ReadFrame, WriteFrame> | |||||
| where | |||||
| ReadFrame: ValueDecode + Debug, | |||||
| WriteFrame: ValueEncode + Debug, | |||||
| { | |||||
| /// Wraps the given `stream` to yield a message channel. | |||||
| pub fn new(stream: TcpStream) -> Self { | |||||
| let (read_half, write_half) = stream.into_split(); | |||||
| Self { | |||||
| reader: ChannelReader { | |||||
| inner: FrameReader::new(read_half), | |||||
| }, | |||||
| writer: FrameWriter::new(write_half), | |||||
| } | |||||
| } | |||||
| // This future sends all the requests from `request_stream` through `writer` | |||||
| // until the stream is finished, then resolves. | |||||
| async fn send<S: Stream<Item = WriteFrame>>( | |||||
| writer: &mut FrameWriter<WriteFrame, OwnedWriteHalf>, | |||||
| send_stream: S, | |||||
| ) -> io::Result<()> { | |||||
| tokio::pin!(send_stream); | |||||
| while let Some(frame) = send_stream.next().await { | |||||
| debug!("Channel: sending frame: {:?}", frame); | |||||
| writer.write(&frame).await?; | |||||
| } | |||||
| Ok(()) | |||||
| } | |||||
| // It would be easier to inline this `select!` call inside `run()`, but that | |||||
| // fails due to some weird, undiagnosed error due to the interaction of | |||||
| // `async_stream::try_stream!`, `select!` and the `?` operator. | |||||
| async fn run_once<S: Future<Output = io::Result<()>>>( | |||||
| send_task: S, | |||||
| reader: &mut ChannelReader<ReadFrame>, | |||||
| ) -> Result<Option<ReadFrame>, ChannelError> { | |||||
| tokio::select! { | |||||
| send_result = send_task => { | |||||
| send_result?; | |||||
| Ok(None) | |||||
| }, | |||||
| read_result = reader.read_strict() => read_result.map(Some), | |||||
| } | |||||
| } | |||||
| /// Attempts to read a single frame from the underlying stream. | |||||
| pub async fn read(&mut self) -> Result<ReadFrame, ChannelError> { | |||||
| self.reader.read_strict().await | |||||
| } | |||||
| /// Attempts to write a single frame to the underlying stream. | |||||
| pub async fn write(&mut self, frame: &WriteFrame) -> io::Result<()> { | |||||
| self.writer.write(frame).await | |||||
| } | |||||
| /// Sends the given stream of frames while receiving frames in return. | |||||
| /// Once `send_stream` is exhausted, shuts down the underlying TCP stream, | |||||
| /// drains incoming frames, then terminates. | |||||
| pub fn run<S: Stream<Item = WriteFrame>>( | |||||
| mut self, | |||||
| send_stream: S, | |||||
| ) -> impl Stream<Item = Result<ReadFrame, ChannelError>> { | |||||
| async_stream::try_stream! { | |||||
| // Drive the main loop: send requests and receive responses. | |||||
| // | |||||
| // We make a big future out of the operation of waiting for requests | |||||
| // to send and from `request_stream` and sending them out through | |||||
| // `self.writer`, that we can then poll repeatedly and concurrently | |||||
| // with polling for responses. This allows us to concurrently write | |||||
| // and read from the underlying `TcpStream` in full duplex mode. | |||||
| { | |||||
| let send_task = Self::send(&mut self.writer, send_stream); | |||||
| tokio::pin!(send_task); | |||||
| while let Some(frame) = | |||||
| Self::run_once(&mut send_task, &mut self.reader).await? { | |||||
| yield frame; | |||||
| } | |||||
| } | |||||
| debug!("Channel: shutting down writer"); | |||||
| self.writer.shutdown().await?; | |||||
| // Drain the receiving end of the connection. | |||||
| while let Some(frame) = self.reader.read().await? { | |||||
| yield frame; | |||||
| } | |||||
| } | |||||
| } | |||||
| } | |||||
| #[cfg(test)] | |||||
| mod tests { | |||||
| use futures::stream::{self, StreamExt}; | |||||
| use tokio::net::{TcpListener, TcpStream}; | |||||
| use tokio::sync::oneshot; | |||||
| use super::Channel; | |||||
| // Enable capturing logs in tests. | |||||
| fn init() { | |||||
| let _ = env_logger::builder().is_test(true).try_init(); | |||||
| } | |||||
| #[tokio::test] | |||||
| async fn read_write() { | |||||
| init(); | |||||
| let listener = TcpListener::bind("localhost:0").await.unwrap(); | |||||
| let address = listener.local_addr().unwrap(); | |||||
| let listener_task = tokio::spawn(async move { | |||||
| let (stream, _) = listener.accept().await.unwrap(); | |||||
| let mut channel = Channel::<u32, u32>::new(stream); | |||||
| assert_eq!(channel.read().await.unwrap(), 1); | |||||
| channel.write(&2).await.unwrap(); | |||||
| }); | |||||
| let stream = TcpStream::connect(address).await.unwrap(); | |||||
| let mut channel = Channel::<u32, u32>::new(stream); | |||||
| channel.write(&1).await.unwrap(); | |||||
| assert_eq!(channel.read().await.unwrap(), 2); | |||||
| listener_task.await.unwrap(); | |||||
| } | |||||
| #[tokio::test] | |||||
| async fn read_eof() { | |||||
| init(); | |||||
| let listener = TcpListener::bind("localhost:0").await.unwrap(); | |||||
| let address = listener.local_addr().unwrap(); | |||||
| let listener_task = tokio::spawn(async move { | |||||
| // Accept the stream and immediately drop/close it. | |||||
| listener.accept().await.unwrap(); | |||||
| }); | |||||
| let stream = TcpStream::connect(address).await.unwrap(); | |||||
| let mut channel = Channel::<u32, u32>::new(stream); | |||||
| assert!(channel.read().await.unwrap_err().is_unexpected_eof()); | |||||
| listener_task.await.unwrap(); | |||||
| } | |||||
| #[tokio::test] | |||||
| async fn open_close() { | |||||
| init(); | |||||
| let listener = TcpListener::bind("localhost:0").await.unwrap(); | |||||
| let address = listener.local_addr().unwrap(); | |||||
| let listener_task = tokio::spawn(async move { | |||||
| let (stream, _) = listener.accept().await.unwrap(); | |||||
| let channel = Channel::<u32, u32>::new(stream); | |||||
| // Wait forever, receive no messages. | |||||
| let inbound = channel.run(stream::pending()); | |||||
| tokio::pin!(inbound); | |||||
| // The server observes an unexpected EOF error when the client | |||||
| // decides to close the channel. | |||||
| let error = inbound.next().await.unwrap().unwrap_err(); | |||||
| assert!(error.is_unexpected_eof()); | |||||
| }); | |||||
| let stream = TcpStream::connect(address).await.unwrap(); | |||||
| let channel = Channel::<u32, u32>::new(stream); | |||||
| // Stop immediately, receive no messages. | |||||
| let inbound = channel.run(stream::empty()); | |||||
| tokio::pin!(inbound); | |||||
| // The channel is closed cleanly from the client's point of view. | |||||
| assert!(inbound.next().await.is_none()); | |||||
| listener_task.await.unwrap(); | |||||
| } | |||||
| #[tokio::test] | |||||
| async fn simple_exchange() { | |||||
| init(); | |||||
| let listener = TcpListener::bind("localhost:0").await.unwrap(); | |||||
| let address = listener.local_addr().unwrap(); | |||||
| let listener_task = tokio::spawn(async move { | |||||
| let (stream, _) = listener.accept().await.unwrap(); | |||||
| let channel = Channel::<u32, u32>::new(stream); | |||||
| let (tx, rx) = oneshot::channel::<u32>(); | |||||
| // Send one message, then wait forever. | |||||
| let outbound = | |||||
| stream::once(async move { rx.await.unwrap() }).chain(stream::pending()); | |||||
| let inbound = channel.run(outbound); | |||||
| tokio::pin!(inbound); | |||||
| // The server receives the client's message. | |||||
| assert_eq!(inbound.next().await.unwrap().unwrap(), 1); | |||||
| // Server responds. | |||||
| tx.send(1001).unwrap(); | |||||
| // The server observes an unexpected EOF error when the client | |||||
| // decides to close the channel. | |||||
| let error = inbound.next().await.unwrap().unwrap_err(); | |||||
| assert!(error.is_unexpected_eof()); | |||||
| }); | |||||
| let stream = TcpStream::connect(address).await.unwrap(); | |||||
| let channel = Channel::<u32, u32>::new(stream); | |||||
| let (tx, rx) = oneshot::channel::<()>(); | |||||
| // Send one message then wait for a reply. If we did not wait, we might | |||||
| // shut down the channel before the server has a chance to respond. | |||||
| let outbound = async_stream::stream! { | |||||
| yield 1; | |||||
| rx.await.unwrap(); | |||||
| }; | |||||
| let inbound = channel.run(outbound); | |||||
| tokio::pin!(inbound); | |||||
| // The client receives the server's message. | |||||
| assert_eq!(inbound.next().await.unwrap().unwrap(), 1001); | |||||
| // Signal to the client that we should shut down. | |||||
| tx.send(()).unwrap(); | |||||
| assert!(inbound.next().await.is_none()); | |||||
| listener_task.await.unwrap(); | |||||
| } | |||||
| } | |||||