Merge pull request #13 from halzy/halzy/10-docs

feat!: Added HalvesStream, example.

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "stream_multiplexer"
-version = "0.4.1"
+version = "0.5.0"
 authors = ["Benjamin Halsted <[email protected]>"]
 edition = "2018"
 license = "MIT OR Apache-2.0"
@@ -17,13 +17,13 @@ byteorder = "1.3"
 bytes = "0.5"
 futures = { version = "0.3", default-features = false, features = ["alloc"] }
 thiserror = "1.0"
-tokio = { version = "0.2", features = ["full"] }
+tokio = { version = "0.2.13", features = ["full"] }
 tokio-util = { version = "0.2", features = ["codec"] }
 tracing = { version = "0.1", features = ["log"] }
 tracing-futures = "0.2"
 
 [dev-dependencies]
 futures = { version = "0.3", default-features = false, features = ["alloc","std"] }
 matches = "0.1"
-tokio = { version = "0.2", features = ["full", "test-util"] }
+tokio = { version = "0.2.13", features = ["full", "test-util"] }
 tracing-subscriber = "0.2"

src/halves_stream.rs

@@ -0,0 +1,62 @@
+use futures::stream::TryStream;
+use futures::Stream;
+use tokio::io::{AsyncRead, AsyncWrite, ReadHalf, WriteHalf};
+use tokio_util::codec::{FramedRead, FramedWrite, LengthDelimitedCodec};
+
+use std::pin::Pin;
+use std::task::Poll;
+
+/// Takes a Stream<Item=AsyncRead + AsyncWrite> and provides a
+/// Stream<Item=( FramedWrite<WriteHalf, LengthDelimitedCodec>, FramedRead<ReadHalf, LengthDelimitedCodec>)>
+#[derive(Debug)]
+pub struct HalvesStream<St> {
+    inner: St,
+    length_field_length: usize,
+}
+
+impl<St> HalvesStream<St> {
+    /// Takes a TcpListener to help own the listener while producing TcpStreams
+    pub fn new(inner: St, length_field_length: usize) -> Self {
+        Self {
+            inner,
+            length_field_length,
+        }
+    }
+}
+
+impl<St> Stream for HalvesStream<St>
+where
+    St: TryStream<Error = std::io::Error> + Unpin,
+    St::Ok: AsyncRead + AsyncWrite,
+{
+    type Item = Result<
+        (
+            FramedWrite<WriteHalf<St::Ok>, LengthDelimitedCodec>,
+            FramedRead<ReadHalf<St::Ok>, LengthDelimitedCodec>,
+        ),
+        St::Error,
+    >;
+    fn poll_next(
+        mut self: Pin<&mut Self>,
+        ctx: &mut std::task::Context,
+    ) -> Poll<Option<Self::Item>> {
+        match futures::ready!(Pin::new(&mut self.inner).try_poll_next(ctx)) {
+            None => None.into(),
+            Some(Err(err)) => Poll::Ready(Some(Err(err))),
+            Some(Ok(stream)) => {
+                let (reader, writer) = tokio::io::split(stream);
+
+                // Wrap the writer in a FramedCodec
+                let framed_write = LengthDelimitedCodec::builder()
+                    .length_field_length(self.length_field_length)
+                    .new_write(writer);
+
+                let framed_read = LengthDelimitedCodec::builder()
+                    .length_field_length(self.length_field_length)
+                    .new_read(reader);
+
+                Poll::Ready(Some(Ok((framed_write, framed_read))))
+            }
+        }
+    }
+}

src/lib.rs

@@ -20,9 +20,121 @@ This crate provides stream multiplexing with channels.
 Channels have their own backpressure that does not affect other channels.
 
 Incoming streams are by default set to channel 0 and can be moved to other channels via `ControlMessage`s.
+
+```rust
+use bytes::Bytes;
+use tokio::net::TcpListener;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::sync::mpsc;
+use stream_multiplexer::{Multiplexer, HalvesStream, ControlMessage, IncomingPacket, OutgoingPacket};
+use futures::stream::StreamExt;
+
+# tokio::runtime::Builder::new().basic_scheduler().enable_all().build().unwrap().block_on(async move {
+// 3 channels of incoming streams, 0 is the channel that new streams join.
+// Backpressure is per channel. Streams can be moved between channels by
+// sending an OutgoingPackt::ChangeChannel message.
+let (channel0_tx, mut channel0_rx) = mpsc::channel(32);
+let (channel1_tx, mut channel1_rx) = mpsc::channel(32);
+let (channel2_tx, mut channel2_rx) = mpsc::channel(32);
+
+// A Stream for outgoing messages.
+let (mut outgoing_tx, outgoing_rx) = mpsc::channel::<OutgoingPacket<Bytes>>(32);
+
+// Construct the multiplexer, giving it the OutgoingPacket stream, and a vector of incoming
+// streams. The backlog controls how much of an internal buffer each WriteHalf (TcpSocket in this example) can have.
+let outgoing_streams_backlog = 128;
+let multiplexer = Multiplexer::new(
+    outgoing_streams_backlog,
+    outgoing_rx,
+    vec![channel0_tx, channel1_tx, channel2_tx],
+);
+
+// Bind to a random port on localhost
+let socket = TcpListener::bind("127.0.0.1:0").await.unwrap();
+
+let local_addr = socket.local_addr().unwrap();
+
+// Use the HalvesStream utility struct to map the stream of new sockets.
+// It will use LengthDelimitedCodec with 2 bytes as the packet size.
+let halves = HalvesStream::new(socket, 2);
+
+// Control channel for shutting down the multiplexer
+let (control_write, control_read) = mpsc::unbounded_channel();
+let mp_joinhandle = tokio::task::spawn(multiplexer.run(halves, control_read));
+
+// Make a test connection:
+let mut client = tokio::net::TcpStream::connect(local_addr).await.unwrap();
+
+// Send 'a message'
+let mut data = Bytes::from("\x00\x09a message");
+client.write_buf(&mut data).await.unwrap();
+client.flush();
+
+// Receive 'a message' on channel 0
+let incoming_packet = channel0_rx.recv().await.unwrap();
+assert_eq!(
+    incoming_packet
+        .value()
+        .expect("should have a value")
+        .as_ref()
+        .unwrap(),
+    &Bytes::from("a message")
+);
+
+// Move the client to channel 1
+outgoing_tx
+    .send(OutgoingPacket::ChangeChannel(vec![incoming_packet.id()], 1))
+    .await
+    .unwrap();
+
+// Send 'a message' again, on channel 1 this time.
+let mut data = Bytes::from("\x00\x09a message");
+client.write_buf(&mut data).await.unwrap();
+client.flush();
+
+// Receive 'a message' on channel 1
+let incoming_packet = channel1_rx.recv().await.unwrap();
+assert_eq!(
+    incoming_packet
+        .value()
+        .expect("should have a value")
+        .as_ref()
+        .unwrap(),
+    &Bytes::from("a message")
+);
+
+// Move the client to channel 2
+outgoing_tx
+    .send(OutgoingPacket::ChangeChannel(vec![incoming_packet.id()], 2))
+    .await
+    .unwrap();
+
+// Send 'a message' again, on channel 2 this time.
+let mut data = Bytes::from("\x00\x09a message");
+client.write_buf(&mut data).await.unwrap();
+client.flush();
+
+// Receive 'a message' on channel 2
+let incoming_packet = channel2_rx.recv().await.unwrap();
+assert_eq!(
+    incoming_packet
+        .value()
+        .expect("should have a value")
+        .as_ref()
+        .unwrap(),
+    &Bytes::from("a message")
+);
+
+// Tell multiplexer te shut down
+control_write.send(ControlMessage::Shutdown).unwrap();
+
+mp_joinhandle.await.unwrap();
+# });
+```
 */
 mod error;
 mod halt;
+mod halves_stream;
 mod id_gen;
 mod multiplexer;
 mod multiplexer_senders;
@@ -32,6 +144,7 @@ mod stream_mover;
 
 pub use error::*;
 use halt::*;
+pub use halves_stream::*;
 pub use id_gen::*;
 pub use multiplexer::*;
 use multiplexer_senders::*;
@@ -45,6 +158,7 @@ type StreamId = usize;
 pub struct IncomingMessage<V> {
     /// Stream Id that the message if for
     pub id: StreamId,
+
     /// Value received from a stream
     pub value: V,
 }

src/multiplexer_senders.rs

@@ -260,7 +260,7 @@ mod tests {
 
     #[tokio::test(basic_scheduler)]
     async fn send_message_and_shutdown() {
-        // crate::tests::init_logging();
+        //crate::tests::init_logging();
 
         // Set up a teest incrementer
         let id_gen = IncrementIdGen::default();