impl recv and poll_recv

Author: estin

Cargo.toml

@@ -22,4 +22,3 @@ 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

@@ -1,4 +1,3 @@
-use futures_util::StreamExt;
 use ntex_redis::{cmd, RedisConnector};
 use std::error::Error;
 
@@ -10,11 +9,11 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let client = RedisConnector::new("127.0.0.1:6379")
         .connect_simple()
         .await?;
-    let mut subscriber = client.stream(cmd::Subscribe("pubsub"))?;
+    let subscriber = client.stream(cmd::Subscribe("pubsub"))?;
 
     ntex::rt::spawn(async move {
         loop {
-            match subscriber.next().await {
+            match subscriber.recv().await {
                 Some(Ok(cmd::SubscribeItem::Subscribed)) => println!("sub: subscribed"),
                 Some(Ok(cmd::SubscribeItem::Message(payload))) => {
                     println!("sub: {:?}", payload)

src/cmd/pubsub.rs

@@ -1,5 +1,6 @@
 use super::{utils, Command, CommandError};
 use ntex::util::Bytes;
+use std::convert::TryFrom;
 
 use crate::codec::{BulkString, Request, Response};
 
@@ -24,6 +25,56 @@ pub enum SubscribeItem {
     Message(Bytes),
 }
 
+impl TryFrom<Response> for SubscribeItem {
+    type Error = CommandError;
+
+    fn try_from(val: Response) -> Result<Self, Self::Error> {
+        let parts = if let Response::Array(ref parts) = val {
+            parts
+        } else {
+            return Err(CommandError::Output("Cannot parse response", val));
+        };
+
+        if parts.len() != 3 {
+            return Err(CommandError::Output(
+                "Subscription message has invalid items number",
+                val,
+            ));
+        }
+
+        let (mtype, payload) = (&parts[0], &parts[2]);
+
+        let mtype = if let Response::Bytes(mtype) = mtype {
+            mtype
+        } else {
+            return Err(CommandError::Output(
+                "Subscription message type unknown",
+                val,
+            ));
+        };
+
+        if mtype == &TYPE_SUBSCRIBE {
+            return Ok(SubscribeItem::Subscribed);
+        }
+
+        if mtype != &TYPE_MESSAGE {
+            return Err(CommandError::Output(
+                "Subscription message type unknown",
+                val,
+            ));
+        }
+
+        if let Response::Bytes(payload) = payload {
+            return Ok(SubscribeItem::Message(payload.clone()));
+        } else {
+            return Err(CommandError::Output(
+                "Subscription message has empty payload",
+                val,
+            ));
+        }
+    }
+}
+
 pub struct SubscribeOutputCommand(pub(crate) Request);
 
 impl Command for SubscribeOutputCommand {
@@ -34,30 +85,7 @@ impl Command for SubscribeOutputCommand {
     }
 
     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)),
-        }
+        SubscribeItem::try_from(val)
     }
 }
 

src/simple.rs

@@ -47,21 +47,16 @@ impl SimpleClient {
     }
 
     /// Execute redis command and stream response
-    pub fn stream<U>(
-        self,
-        cmd: U,
-    ) -> Result<impl Stream<Item = Result<U::Output, CommandError>>, CommandError>
+    pub fn stream<U>(self, cmd: U) -> Result<RedisStream<U>, CommandError>
     where
         U: Command,
     {
         self.io.encode(cmd.to_request(), &Codec)?;
 
-        let rs: RedisStream<U> = RedisStream {
+        Ok(RedisStream {
             io: self.io,
             _cmd: std::marker::PhantomData,
-        };
-
-        Ok(rs)
+        })
     }
 
     pub(crate) fn into_inner(self) -> IoBoxed {
@@ -78,6 +73,23 @@ 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
+    }
+
+    /// 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>>> {
         match ready!(self.io.poll_recv(&Codec, cx)) {
             Ok(item) => match item.into_result() {
                 Ok(result) => Poll::Ready(Some(U::to_output(result))),

tests/test_redis.rs

@@ -1,4 +1,3 @@
-use futures_util::StreamExt;
 use ntex::util::{Bytes, HashMap};
 use ntex_redis::{cmd, Client, RedisConnector};
 use rand::{distributions::Alphanumeric, thread_rng, Rng};
@@ -211,8 +210,8 @@ async fn test_pubsub() {
         .await
         .unwrap();
 
-    let mut stream = subscriber.stream(cmd::Subscribe(&key)).unwrap();
-    let message = stream.next().await;
+    let stream = subscriber.stream(cmd::Subscribe(&key)).unwrap();
+    let message = stream.recv().await;
     assert_eq!(message.unwrap().unwrap(), cmd::SubscribeItem::Subscribed);
 
     let publisher = connect().await;
@@ -222,7 +221,7 @@ async fn test_pubsub() {
     assert_eq!(result, 1);
 
     // sub
-    let message = stream.next().await;
+    let message = stream.recv().await;
     assert_eq!(
         message.unwrap().unwrap(),
         cmd::SubscribeItem::Message(Bytes::from_static(b"1"))