pubsub initial support

Author: estin

Cargo.toml

@@ -22,3 +22,4 @@ derive_more = "0.99"
 rand = "0.8"
 env_logger = "0.9"
 ntex = { version = "0.5", features = ["tokio"] }
+futures-util = "0.3.21"

examples/pubsub.rs

@@ -0,0 +1,41 @@
+use futures_util::StreamExt;
+use ntex_redis::{cmd, RedisConnector};
+use std::error::Error;
+
+#[ntex::main]
+async fn main() -> Result<(), Box<dyn Error>> {
+    env_logger::init();
+
+    // subscribe
+    let client = RedisConnector::new("127.0.0.1:6379")
+        .connect_simple()
+        .await?;
+    let mut subscriber = client.stream(cmd::Subscribe("pubsub"))?;
+
+    ntex::rt::spawn(async move {
+        loop {
+            match subscriber.next().await {
+                Some(Ok(cmd::SubscribeItem::Subscribed)) => println!("sub: subscribed"),
+                Some(Ok(cmd::SubscribeItem::Message(payload))) => {
+                    println!("sub: {:?}", payload)
+                }
+                Some(Err(e)) => {
+                    println!("sub: {}", e);
+                    return;
+                }
+                _ => unreachable!(),
+            }
+        }
+    });
+
+    // publish
+    let redis = RedisConnector::new("127.0.0.1:6379").connect().await?;
+
+    for i in 0..5 {
+        let value = i.to_string();
+        println!("pub: {}", value);
+        redis.exec(cmd::Publish("pubsub", &value)).await?;
+    }
+
+    Ok(())
+}

src/cmd/mod.rs

@@ -9,6 +9,7 @@ mod connection;
 mod hashes;
 mod keys;
 mod lists;
+mod pubsub;
 mod strings;
 mod utils;
 
@@ -17,6 +18,7 @@ pub use self::connection::{Ping, Select};
 pub use self::hashes::{HDel, HGet, HGetAll, HIncrBy, HLen, HSet};
 pub use self::keys::{Del, Exists, Expire, ExpireAt, Keys, Ttl, TtlResult};
 pub use self::lists::{LIndex, LPop, LPush, RPop, RPush};
+pub use self::pubsub::{Publish, Subscribe, SubscribeItem};
 pub use self::strings::{Get, IncrBy, Set};
 
 /// Trait implemented by types that can be used as redis commands

src/cmd/pubsub.rs

@@ -0,0 +1,73 @@
+use super::{utils, Command, CommandError};
+use ntex::util::Bytes;
+
+use crate::codec::{BulkString, Request, Response};
+
+const TYPE_SUBSCRIBE: Bytes = Bytes::from_static(b"subscribe");
+const TYPE_MESSAGE: Bytes = Bytes::from_static(b"message");
+
+/// Publish redis command
+pub fn Publish<T, V>(key: T, value: V) -> utils::IntOutputCommand
+where
+    BulkString: From<T> + From<V>,
+{
+    utils::IntOutputCommand(Request::Array(vec![
+        Request::from_static("PUBLISH"),
+        Request::BulkString(key.into()),
+        Request::BulkString(value.into()),
+    ]))
+}
+
+#[derive(Debug, PartialEq)]
+pub enum SubscribeItem {
+    Subscribed,
+    Message(Bytes),
+}
+
+pub struct SubscribeOutputCommand(pub(crate) Request);
+
+impl Command for SubscribeOutputCommand {
+    type Output = SubscribeItem;
+
+    fn to_request(self) -> Request {
+        self.0
+    }
+
+    fn to_output(val: Response) -> Result<Self::Output, CommandError> {
+        match val {
+            Response::Array(ref v) => match v.get(0) {
+                Some(Response::Bytes(t)) => {
+                    if t == &TYPE_SUBSCRIBE {
+                        return Ok(SubscribeItem::Subscribed);
+                    }
+                    if t == &TYPE_MESSAGE {
+                        return match v.get(2) {
+                            Some(payload) => match payload {
+                                Response::Bytes(m) => Ok(SubscribeItem::Message(m.clone())),
+                                _ => {
+                                    Err(CommandError::Output("Cannot parse message payload", val))
+                                }
+                            },
+                            _ => Err(CommandError::Output("Empty messsage payload", val)),
+                        };
+                    }
+
+                    Err(CommandError::Output("Unknown message type", val))
+                }
+                _ => Err(CommandError::Output("Cannot parse message type", val)),
+            },
+            _ => Err(CommandError::Output("Cannot parse response", val)),
+        }
+    }
+}
+
+/// Subscribe redis command
+pub fn Subscribe<T>(key: T) -> SubscribeOutputCommand
+where
+    BulkString: From<T>,
+{
+    SubscribeOutputCommand(Request::Array(vec![
+        Request::from_static("SUBSCRIBE"),
+        Request::BulkString(key.into()),
+    ]))
+}

src/simple.rs

@@ -1,10 +1,10 @@
-use std::task::Poll;
-
-use ntex::{io::IoBoxed, io::RecvError, util::poll_fn, util::ready};
+use std::pin::Pin;
+use std::task::{Context, Poll};
 
 use super::cmd::Command;
 use super::codec::Codec;
 use super::errors::{CommandError, Error};
+use ntex::{io::IoBoxed, io::RecvError, util::poll_fn, util::ready, util::Stream};
 
 /// Redis client
 pub struct SimpleClient {
@@ -46,7 +46,59 @@ impl SimpleClient {
         .await
     }
 
+    /// Execute redis command and stream response
+    pub fn stream<U>(
+        self,
+        cmd: U,
+    ) -> Result<impl Stream<Item = Result<U::Output, CommandError>>, CommandError>
+    where
+        U: Command,
+    {
+        self.io.encode(cmd.to_request(), &Codec)?;
+
+        let rs: RedisStream<U> = RedisStream {
+            io: self.io,
+            _cmd: std::marker::PhantomData,
+        };
+
+        Ok(rs)
+    }
+
     pub(crate) fn into_inner(self) -> IoBoxed {
         self.io
     }
 }
+
+pub struct RedisStream<U: Command> {
+    io: IoBoxed,
+    _cmd: std::marker::PhantomData<U>,
+}
+
+impl<U: Command> Stream for RedisStream<U> {
+    type Item = Result<U::Output, CommandError>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match ready!(self.io.poll_recv(&Codec, cx)) {
+            Ok(item) => match item.into_result() {
+                Ok(result) => Poll::Ready(Some(U::to_output(result))),
+                Err(err) => Poll::Ready(Some(Err(CommandError::Error(err)))),
+            },
+            Err(RecvError::KeepAlive) | Err(RecvError::Stop) => {
+                unreachable!()
+            }
+            Err(RecvError::WriteBackpressure) => {
+                if let Err(err) = ready!(self.io.poll_flush(cx, false))
+                    .map_err(|e| CommandError::Protocol(Error::PeerGone(Some(e))))
+                {
+                    Poll::Ready(Some(Err(err)))
+                } else {
+                    Poll::Pending
+                }
+            }
+            Err(RecvError::Decoder(err)) => Poll::Ready(Some(Err(CommandError::Protocol(err)))),
+            Err(RecvError::PeerGone(err)) => {
+                Poll::Ready(Some(Err(CommandError::Protocol(Error::PeerGone(err)))))
+            }
+        }
+    }
+}

tests/test_redis.rs

@@ -1,4 +1,5 @@
-use ntex::util::HashMap;
+use futures_util::StreamExt;
+use ntex::util::{Bytes, HashMap};
 use ntex_redis::{cmd, Client, RedisConnector};
 use rand::{distributions::Alphanumeric, thread_rng, Rng};
 use std::time::{Duration, SystemTime};
@@ -200,3 +201,30 @@ async fn test_connection() {
     let result = redis.exec(cmd::Select(1)).await.unwrap();
     assert!(result);
 }
+
+#[ntex::test]
+async fn test_pubsub() {
+    let key = new_key();
+
+    let subscriber = RedisConnector::new("127.0.0.1:6379")
+        .connect_simple()
+        .await
+        .unwrap();
+
+    let mut stream = subscriber.stream(cmd::Subscribe(&key)).unwrap();
+    let message = stream.next().await;
+    assert_eq!(message.unwrap().unwrap(), cmd::SubscribeItem::Subscribed);
+
+    let publisher = connect().await;
+
+    // pub
+    let result = publisher.exec(cmd::Publish(&key, "1")).await.unwrap();
+    assert_eq!(result, 1);
+
+    // sub
+    let message = stream.next().await;
+    assert_eq!(
+        message.unwrap().unwrap(),
+        cmd::SubscribeItem::Message(Bytes::from_static(b"1"))
+    );
+}