10 changed files with 273 additions and 522 deletions
Split 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(); | |||
| } | |||
| } | |||