don't consume connection, unsubscribe command

Author: estin

examples/pubsub.rs

@@ -1,23 +1,30 @@
 use ntex_redis::{cmd, RedisConnector};
 use std::error::Error;
+use std::rc::Rc;
 
 #[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 subscriber = client.stream(cmd::Subscribe("pubsub"))?;
+    // subscriber
+    let client = Rc::new(
+        RedisConnector::new("127.0.0.1:6379")
+            .connect_simple()
+            .await?,
+    );
+
+    let client_clone = client.clone();
 
     ntex::rt::spawn(async move {
+        let subscriber = client_clone.stream(cmd::Subscribe("pubsub")).unwrap();
+
         loop {
             match subscriber.recv().await {
                 Some(Ok(cmd::SubscribeItem::Subscribed)) => println!("sub: subscribed"),
                 Some(Ok(cmd::SubscribeItem::Message(payload))) => {
                     println!("sub: {:?}", payload)
                 }
+                Some(Ok(cmd::SubscribeItem::UnSubscribed)) => println!("sub: unsubscribed"),
                 Some(Err(e)) => {
                     println!("sub: {}", e);
                     return;
@@ -36,5 +43,11 @@ async fn main() -> Result<(), Box<dyn Error>> {
         redis.exec(cmd::Publish("pubsub", &value)).await?;
     }
 
+    // unsubscribe
+    client.send(cmd::UnSubscribe("pubsub"))?;
+
+    // allow to subscriber recv unsubscribe message
+    ntex::time::sleep(ntex::time::Millis(10)).await;
+
     Ok(())
 }

src/cmd/mod.rs

@@ -18,7 +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::pubsub::{Publish, Subscribe, SubscribeItem, UnSubscribe};
 pub use self::strings::{Get, IncrBy, Set};
 
 /// Trait implemented by types that can be used as redis commands

src/cmd/pubsub.rs

@@ -5,6 +5,7 @@ use std::convert::TryFrom;
 use crate::codec::{BulkString, Request, Response};
 
 const TYPE_SUBSCRIBE: Bytes = Bytes::from_static(b"subscribe");
+const TYPE_UNSUBSCRIBE: Bytes = Bytes::from_static(b"unsubscribe");
 const TYPE_MESSAGE: Bytes = Bytes::from_static(b"message");
 
 /// Publish redis command
@@ -22,6 +23,7 @@ where
 #[derive(Debug, PartialEq)]
 pub enum SubscribeItem {
     Subscribed,
+    UnSubscribed,
     Message(Bytes),
 }
 
@@ -57,6 +59,10 @@ impl TryFrom<Response> for SubscribeItem {
             return Ok(SubscribeItem::Subscribed);
         }
 
+        if mtype == &TYPE_UNSUBSCRIBE {
+            return Ok(SubscribeItem::UnSubscribed);
+        }
+
         if mtype != &TYPE_MESSAGE {
             return Err(CommandError::Output(
                 "Subscription message type unknown",
@@ -99,3 +105,14 @@ where
         Request::BulkString(key.into()),
     ]))
 }
+
+/// Unsubscribe redis command
+pub fn UnSubscribe<T>(key: T) -> SubscribeOutputCommand
+where
+    BulkString: From<T>,
+{
+    SubscribeOutputCommand(Request::Array(vec![
+        Request::from_static("UNSUBSCRIBE"),
+        Request::BulkString(key.into()),
+    ]))
+}

src/simple.rs

@@ -17,76 +17,45 @@ impl SimpleClient {
         SimpleClient { io }
     }
 
-    /// Execute redis command
+    /// Execute redis command and wait result
     pub async fn exec<U>(&self, cmd: U) -> Result<U::Output, CommandError>
     where
         U: Command,
     {
-        self.io.encode(cmd.to_request(), &Codec)?;
-
-        poll_fn(|cx| loop {
-            return match ready!(self.io.poll_recv(&Codec, cx)) {
-                Ok(item) => Poll::Ready(U::to_output(
-                    item.into_result().map_err(CommandError::Error)?,
-                )),
-                Err(RecvError::KeepAlive) | Err(RecvError::Stop) => {
-                    unreachable!()
-                }
-                Err(RecvError::WriteBackpressure) => {
-                    ready!(self.io.poll_flush(cx, false))
-                        .map_err(|e| CommandError::Protocol(Error::PeerGone(Some(e))))?;
-                    continue;
-                }
-                Err(RecvError::Decoder(err)) => Poll::Ready(Err(CommandError::Protocol(err))),
-                Err(RecvError::PeerGone(err)) => {
-                    Poll::Ready(Err(CommandError::Protocol(Error::PeerGone(err))))
-                }
-            };
-        })
-        .await
+        self.send(cmd)?;
+        self.recv::<U>().await.unwrap()
     }
 
-    /// Execute redis command and stream response
-    pub fn stream<U>(self, cmd: U) -> Result<RedisStream<U>, CommandError>
+    /// Send redis command
+    pub fn send<U>(&self, cmd: U) -> Result<(), CommandError>
     where
         U: Command,
     {
         self.io.encode(cmd.to_request(), &Codec)?;
+        Ok(())
+    }
 
-        Ok(RedisStream {
-            io: self.io,
+    /// Execute redis command and act with output as stream
+    pub fn stream<U>(&self, cmd: U) -> Result<OutputStream<U>, CommandError>
+    where
+        U: Command,
+    {
+        self.send(cmd)?;
+        Ok(OutputStream {
+            client: self,
             _cmd: std::marker::PhantomData,
         })
     }
 
     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>> {
-        self.poll_recv(cx)
-    }
-}
-
-impl<U: Command> RedisStream<U> {
-    /// Attempt to pull out the next value of this stream.
-    pub async fn recv(&self) -> Option<Result<U::Output, CommandError>> {
-        poll_fn(|cx| self.poll_recv(cx)).await
+    async fn recv<U: Command>(&self) -> Option<Result<U::Output, CommandError>> {
+        poll_fn(|cx| self.poll_recv::<U>(cx)).await
     }
 
-    /// Attempt to pull out the next value of this stream, registering
-    /// the current task for wakeup if the value is not yet available,
-    /// and returning None if the payload is exhausted.
-    pub fn poll_recv(
+    fn poll_recv<U: Command>(
         &self,
         cx: &mut Context<'_>,
     ) -> Poll<Option<Result<U::Output, CommandError>>> {
@@ -114,3 +83,33 @@ impl<U: Command> RedisStream<U> {
         }
     }
 }
+
+pub struct OutputStream<'a, U> {
+    client: &'a SimpleClient,
+    _cmd: std::marker::PhantomData<U>,
+}
+
+impl<'a, U: Command> Stream for OutputStream<'a, U> {
+    type Item = Result<U::Output, CommandError>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.client.poll_recv::<U>(cx)
+    }
+}
+
+impl<'a, U: Command> OutputStream<'a, U> {
+    /// Attempt to pull out the next value of this stream.
+    pub async fn recv(&self) -> Option<Result<U::Output, CommandError>> {
+        poll_fn(|cx| self.client.poll_recv::<U>(cx)).await
+    }
+
+    /// Attempt to pull out the next value of this stream, registering
+    /// the current task for wakeup if the value is not yet available,
+    /// and returning None if the payload is exhausted.
+    pub fn poll_recv(
+        &self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<U::Output, CommandError>>> {
+        self.client.poll_recv::<U>(cx)
+    }
+}

tests/test_redis.rs

@@ -211,6 +211,7 @@ async fn test_pubsub() {
         .unwrap();
 
     let stream = subscriber.stream(cmd::Subscribe(&key)).unwrap();
+
     let message = stream.recv().await;
     assert_eq!(message.unwrap().unwrap(), cmd::SubscribeItem::Subscribed);
 
@@ -226,4 +227,10 @@ async fn test_pubsub() {
         message.unwrap().unwrap(),
         cmd::SubscribeItem::Message(Bytes::from_static(b"1"))
     );
+
+    // unsub
+    subscriber.send(cmd::UnSubscribe(&key)).unwrap();
+
+    let message = stream.recv().await;
+    assert_eq!(message.unwrap().unwrap(), cmd::SubscribeItem::UnSubscribed);
 }