From 3fbd9ddc429607ba9a072781b6701dc6f1a30030 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Oliveira?= Date: Sun, 5 Apr 2020 18:33:21 +0100 Subject: [PATCH] add pub sub client implementation with examples (#22) * add pub sub client implementation with examples * replace subscribed_channels list Vec with HashSet to avoid duplicates * update Subscriber to use async-stream instead of manual Stream impl * revert update to error handling server.rs, as #21 handles it * remove uneeded recursion limit extension --- Cargo.lock | 22 ++++ Cargo.toml | 1 + examples/pub.rs | 31 +++++- examples/sub.rs | 37 ++++++- src/client.rs | 251 +++++++++++++++++++++++++++++++++++++++++-- src/cmd/publish.rs | 13 ++- src/cmd/subscribe.rs | 30 ++++-- 7 files changed, 363 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 23dec1a..3d77455 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -29,6 +29,25 @@ name = "arc-swap" version = "0.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "async-stream" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "async-stream-impl 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)", + "futures-core 0.3.4 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "async-stream-impl" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "proc-macro2 1.0.10 (registry+https://github.com/rust-lang/crates.io-index)", + "quote 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)", + "syn 1.0.17 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "atoi" version = "0.3.2" @@ -302,6 +321,7 @@ dependencies = [ name = "mini-redis" version = "0.1.0" dependencies = [ + "async-stream 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)", "atoi 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)", "bytes 0.5.4 (registry+https://github.com/rust-lang/crates.io-index)", "clap 3.0.0-beta.1 (git+https://github.com/clap-rs/clap/)", @@ -1011,6 +1031,8 @@ dependencies = [ "checksum ansi_term 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)" = "ee49baf6cb617b853aa8d93bf420db2383fab46d314482ca2803b40d5fde979b" "checksum ansi_term 0.12.1 (registry+https://github.com/rust-lang/crates.io-index)" = "d52a9bb7ec0cf484c551830a7ce27bd20d67eac647e1befb56b0be4ee39a55d2" "checksum arc-swap 0.4.5 (registry+https://github.com/rust-lang/crates.io-index)" = "d663a8e9a99154b5fb793032533f6328da35e23aac63d5c152279aa8ba356825" +"checksum async-stream 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "22068c0c19514942eefcfd4daf8976ef1aad84e61539f95cd200c35202f80af5" +"checksum async-stream-impl 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "25f9db3b38af870bf7e5cc649167533b493928e50744e2c30ae350230b414670" "checksum atoi 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)" = "e0afb7287b68575f5ca0e5c7e40191cbd4be59d325781f46faa603e176eaef47" "checksum atty 0.2.14 (registry+https://github.com/rust-lang/crates.io-index)" = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" "checksum autocfg 1.0.0 (registry+https://github.com/rust-lang/crates.io-index)" = "f8aac770f1885fd7e387acedd76065302551364496e46b3dd00860b2f8359b9d" diff --git a/Cargo.toml b/Cargo.toml index a8d8009..202689c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ tokio = { git = "https://github.com/tokio-rs/tokio", features = ["full"] } tracing = "0.1.13" tracing-futures = { version = "0.2.3", features = ["tokio"] } tracing-subscriber = "0.2.2" +async-stream = "0.2.1" [dev-dependencies] # Enable test-utilities in dev mode only. This is mostly for tests. diff --git a/examples/pub.rs b/examples/pub.rs index 607aa27..0a5d05d 100644 --- a/examples/pub.rs +++ b/examples/pub.rs @@ -1,4 +1,31 @@ +//! Publish to a redis channel example. +//! +//! A simple client that connects to a mini-redis server, and +//! publishes a message on `foo` channel +//! +//! You can test this out by running: +//! +//! cargo run --bin server +//! +//! Then in another terminal run: +//! +//! cargo run --example sub +//! +//! And then in another terminal run: +//! +//! cargo run --example pub + +#![warn(rust_2018_idioms)] + +use mini_redis::{client, Result}; + #[tokio::main] -async fn main() { - unimplemented!(); +async fn main() -> Result<()> { + // Open a connection to the mini-redis address. + let mut client = client::connect("127.0.0.1:6379").await?; + + // publish message `bar` on channel foo + client.publish("foo", "bar".into()).await?; + + Ok(()) } diff --git a/examples/sub.rs b/examples/sub.rs index eda175c..ccd3cf9 100644 --- a/examples/sub.rs +++ b/examples/sub.rs @@ -1,6 +1,37 @@ -/// Subscribe to a redis channel +//! Subscribe to a redis channel example. +//! +//! A simple client that connects to a mini-redis server, subscribes to "foo" and "bar" channels +//! and awaits messages published on those channels +//! +//! You can test this out by running: +//! +//! cargo run --bin server +//! +//! Then in another terminal run: +//! +//! cargo run --example sub +//! +//! And then in another terminal run: +//! +//! cargo run --example pub + +#![warn(rust_2018_idioms)] + +use mini_redis::{client, Result}; #[tokio::main] -async fn main() { - unimplemented!(); +pub async fn main() -> Result<()> { + // Open a connection to the mini-redis address. + let client = client::connect("127.0.0.1:6379").await?; + + + // subscribe to channel foo + let mut subscriber = client.subscribe(vec!["foo".into()]).await?; + + // await messages on channel foo + let msg = subscriber.next_message().await? ; + println!("got message from the channel: {}; message = {:?}", msg.channel, msg.content); + + + Ok(()) } diff --git a/src/client.rs b/src/client.rs index 0568976..812e896 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,11 +1,15 @@ +use crate::cmd::{Get, Publish, Set, Subscribe, Unsubscribe}; use crate::{Connection, Frame}; -use crate::cmd::{Get, Set}; use bytes::Bytes; use std::io::{Error, ErrorKind}; +use std::iter::FromIterator; +use std::collections::HashSet; use std::time::Duration; use tokio::net::{TcpStream, ToSocketAddrs}; +use tokio::stream::Stream; use tracing::{debug, instrument}; +use async_stream::stream; /// Mini asynchronous Redis client pub struct Client { @@ -47,7 +51,33 @@ impl Client { key: key.to_string(), value: value, expire: None, - }).await + }) + .await + } + + /// publish `message` on the `channel` + #[instrument(skip(self))] + pub async fn publish(&mut self, channel: &str, message: Bytes) -> crate::Result { + self.publish_cmd(Publish { + channel: channel.to_string(), + message: message, + }) + .await + } + + /// subscribe to the list of channels + /// when client sends a `SUBSCRIBE` command, server's handle for client enters a mode where only + /// `SUBSCRIBE` and `UNSUBSCRIBE` commands are allowed, so we consume client and return Subscribe type + /// which only allows `SUBSCRIBE` and `UNSUBSCRIBE` commands + #[instrument(skip(self))] + pub async fn subscribe(mut self, channels: Vec) -> crate::Result { + let channels = self.subscribe_cmd(Subscribe { channels: channels }).await?; + let subscribed_channels = HashSet::from_iter(channels); + + Ok(Subscriber { + conn: self.conn, + subscribed_channels, + }) } /// Set the value of a key to `value`. The value expires after `expiration`. @@ -62,7 +92,8 @@ impl Client { key: key.to_string(), value: value.into(), expire: Some(expiration), - }).await + }) + .await } async fn set_cmd(&mut self, cmd: Set) -> crate::Result<()> { @@ -81,6 +112,52 @@ impl Client { } } + async fn publish_cmd(&mut self, cmd: Publish) -> crate::Result { + // Convert the `Publish` command into a frame + let frame = cmd.into_frame(); + + debug!(request = ?frame); + + // Write the frame to the socket + self.conn.write_frame(&frame).await?; + + // Read the response + match self.read_response().await? { + Frame::Integer(response) => Ok(response), + frame => Err(frame.to_error()), + } + } + + async fn subscribe_cmd(&mut self, cmd: Subscribe) -> crate::Result> { + // Convert the `Subscribe` command into a frame + let channels = cmd.channels.clone(); + let frame = cmd.into_frame(); + + debug!(request = ?frame); + + // Write the frame to the socket + self.conn.write_frame(&frame).await?; + + // Read the response + for channel in &channels { + let response = self.read_response().await?; + match response { + Frame::Array(ref frame) => match frame.as_slice() { + [subscribe, schannel] + if subscribe.to_string() == "subscribe" + && &schannel.to_string() == channel => + { + () + } + _ => return Err(response.to_error()), + }, + frame => return Err(frame.to_error()), + }; + } + + Ok(channels) + } + /// Reads a response frame from the socket. If an `Error` frame is read, it /// is converted to `Err`. async fn read_response(&mut self) -> crate::Result { @@ -89,20 +166,176 @@ impl Client { debug!(?response); match response { - Some(Frame::Error(msg)) => { - Err(msg.into()) - } + Some(Frame::Error(msg)) => Err(msg.into()), Some(frame) => Ok(frame), None => { // Receiving `None` here indicates the server has closed the // connection without sending a frame. This is unexpected and is // represented as a "connection reset by peer" error. - let err = Error::new( - ErrorKind::ConnectionReset, - "connection reset by server"); + let err = Error::new(ErrorKind::ConnectionReset, "connection reset by server"); Err(err.into()) } } } } + +pub struct Subscriber { + conn: Connection, + subscribed_channels: HashSet, +} + +impl Subscriber { + + /// await for next message published on the subscribed channels + pub async fn next_message(&mut self) -> crate::Result { + match self.receive_message().await { + Some(message) => message, + None => { + // Receiving `None` here indicates the server has closed the + // connection without sending a frame. This is unexpected and is + // represented as a "connection reset by peer" error. + let err = Error::new(ErrorKind::ConnectionReset, "connection reset by server"); + + Err(err.into()) + } + } + } + + /// Convert the subscriber into a Stream + /// yielding new messages published on subscribed channels + pub fn into_stream(mut self) -> impl Stream> { + stream! { + while let Some(message) = self.receive_message().await { + yield message; + } + } + } + + /// Subscribe to a list of new channels + #[instrument(skip(self))] + pub async fn subscribe(&mut self, channels: Vec) -> crate::Result<()> { + let cmd = Subscribe { channels: channels }; + + let channels = cmd.channels.clone(); + let frame = cmd.into_frame(); + + debug!(request = ?frame); + + // Write the frame to the socket + self.conn.write_frame(&frame).await?; + + // Read the response + for channel in &channels { + let response = self.read_response().await?; + match response { + Frame::Array(ref frame) => match frame.as_slice() { + [subscribe, schannel] + if &subscribe.to_string() == "subscribe" + && &schannel.to_string() == channel => + { + () + } + _ => return Err(response.to_error()), + }, + frame => return Err(frame.to_error()), + }; + } + + self.subscribed_channels.extend(channels); + + Ok(()) + } + + /// Unsubscribe to a list of new channels + #[instrument(skip(self))] + pub async fn unsubscribe(&mut self, channels: Vec) -> crate::Result<()> { + let cmd = Unsubscribe { channels: channels }; + + let mut channels = cmd.channels.clone(); + let frame = cmd.into_frame(); + + debug!(request = ?frame); + + // Write the frame to the socket + self.conn.write_frame(&frame).await?; + + // if the input channel list is empty, server acknowledges as unsubscribing + // from all subscribed channels, so we assert that the unsubscribe list received + // matches the client subscribed one + if channels.is_empty() { + channels = Vec::from_iter(self.subscribed_channels.clone()); + } + + // Read the response + for channel in &channels { + let response = self.read_response().await?; + match response { + Frame::Array(ref frame) => match frame.as_slice() { + [unsubscribe, uchannel] + if &unsubscribe.to_string() == "unsubscribe" + && &uchannel.to_string() == channel => + { + self.subscribed_channels.remove(&uchannel.to_string()); + } + _ => return Err(response.to_error()), + }, + frame => return Err(frame.to_error()), + }; + } + + Ok(()) + } + + /// Receives a frame published from server on socket and convert it to a `Message` + /// if frame is not `Frame::Array` with proper message structure return Err + async fn receive_message(&mut self) -> Option> { + match self.conn.read_frame().await { + Ok(None) => None, + Err(err) => Some(Err(err.into())), + Ok(Some(mframe)) => { + debug!(?mframe); + match mframe { + Frame::Array(ref frame) => match frame.as_slice() { + [message, channel, content] if &message.to_string() == "message" => { + Some(Ok(Message { + channel: channel.to_string(), + content: Bytes::from(content.to_string()), + })) + } + _ => Some(Err(mframe.to_error())), + }, + frame => Some(Err(frame.to_error())), + } + } + } + } + + /// Reads a response frame to a command from the socket. If an `Error` frame is read, it + /// is converted to `Err`. + async fn read_response(&mut self) -> crate::Result { + let response = self.conn.read_frame().await?; + + debug!(?response); + + match response { + Some(Frame::Error(msg)) => Err(msg.into()), + Some(frame) => Ok(frame), + None => { + // Receiving `None` here indicates the server has closed the + // connection without sending a frame. This is unexpected and is + // represented as a "connection reset by peer" error. + let err = Error::new(ErrorKind::ConnectionReset, "connection reset by server"); + + Err(err.into()) + } + } + } +} + +/// A message received on a subscribed channel +#[derive(Debug, Clone)] +pub struct Message { + pub channel: String, + pub content: Bytes, +} diff --git a/src/cmd/publish.rs b/src/cmd/publish.rs index 7e937f2..dc13a7e 100644 --- a/src/cmd/publish.rs +++ b/src/cmd/publish.rs @@ -4,8 +4,8 @@ use bytes::Bytes; #[derive(Debug)] pub struct Publish { - channel: String, - message: Bytes, + pub(crate) channel: String, + pub(crate) message: Bytes, } impl Publish { @@ -24,4 +24,13 @@ impl Publish { dst.write_frame(&response).await?; Ok(()) } + + pub(crate) fn into_frame(self) -> Frame { + let mut frame = Frame::array(); + frame.push_bulk(Bytes::from("publish".as_bytes())); + frame.push_bulk(Bytes::from(self.channel.into_bytes())); + frame.push_bulk(self.message); + + frame + } } diff --git a/src/cmd/subscribe.rs b/src/cmd/subscribe.rs index d97202e..7e6cdef 100644 --- a/src/cmd/subscribe.rs +++ b/src/cmd/subscribe.rs @@ -1,5 +1,5 @@ -use crate::{Command, Connection, Db, Frame, Shutdown}; use crate::cmd::{Parse, ParseError}; +use crate::{Command, Connection, Db, Frame, Shutdown}; use bytes::Bytes; use tokio::select; @@ -7,12 +7,12 @@ use tokio::stream::{StreamExt, StreamMap}; #[derive(Debug)] pub struct Subscribe { - channels: Vec, + pub(crate) channels: Vec, } -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct Unsubscribe { - channels: Vec, + pub(crate) channels: Vec, } impl Subscribe { @@ -98,8 +98,8 @@ impl Subscribe { res = dst.read_frame() => { let frame = match res? { Some(frame) => frame, - // How to handle remote client closing write half - None => unimplemented!(), + // How to handle remote client closing write half? + None => return Ok(()) }; // A command has been received from the client. @@ -147,6 +147,15 @@ impl Subscribe { }; } } + + pub(crate) fn into_frame(self) -> Frame { + let mut frame = Frame::array(); + frame.push_bulk(Bytes::from("subscribe".as_bytes())); + for channel in self.channels { + frame.push_bulk(Bytes::from(channel.into_bytes())); + } + frame + } } impl Unsubscribe { @@ -166,4 +175,13 @@ impl Unsubscribe { Ok(Unsubscribe { channels }) } + + pub(crate) fn into_frame(self) -> Frame { + let mut frame = Frame::array(); + frame.push_bulk(Bytes::from("unsubscribe".as_bytes())); + for channel in self.channels { + frame.push_bulk(Bytes::from(channel.into_bytes())); + } + frame + } }