1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use std::pin::Pin;
use std::{future::poll_fn, task::Context, task::Poll};

use super::cmd::{commands::PubSubCommand, commands::SubscribeOutputCommand, Command};
use super::codec::Codec;
use super::errors::{CommandError, Error};
use ntex::{io::IoBoxed, io::RecvError, util::ready, util::Stream};

/// Redis client
pub struct SimpleClient {
    io: IoBoxed,
}

impl SimpleClient {
    /// Create new simple client
    pub(crate) fn new(io: IoBoxed) -> Self {
        SimpleClient { io }
    }

    /// Execute redis command and wait result
    pub async fn exec<U>(&self, cmd: U) -> Result<U::Output, CommandError>
    where
        U: Command,
    {
        self.send(cmd)?;
        loop {
            if let Some(result) = self.recv::<U>().await {
                return result;
            }
        }
    }

    /// Send redis command
    pub fn send<U>(&self, cmd: U) -> Result<(), CommandError>
    where
        U: Command,
    {
        self.io.encode(cmd.to_request(), &Codec)?;
        Ok(())
    }

    /// Execute redis SUBSCRIBE command and act with output as stream
    pub fn subscribe(
        self,
        cmd: SubscribeOutputCommand,
    ) -> Result<SubscriptionClient<SubscribeOutputCommand>, CommandError> {
        self.send(cmd)?;
        Ok(SubscriptionClient {
            client: self,
            _cmd: std::marker::PhantomData,
        })
    }

    pub(crate) fn into_inner(self) -> IoBoxed {
        self.io
    }

    async fn recv<U: Command>(&self) -> Option<Result<U::Output, CommandError>> {
        poll_fn(|cx| self.poll_recv::<U>(cx)).await
    }

    fn poll_recv<U: Command>(
        &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))),
                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)))))
            }
        }
    }
}

/// Redis pubsub client to receive push messages
pub struct SubscriptionClient<U: Command + PubSubCommand> {
    client: SimpleClient,
    _cmd: std::marker::PhantomData<U>,
}

impl<U: Command + PubSubCommand> Stream for SubscriptionClient<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 + PubSubCommand> SubscriptionClient<U> {
    /// Get client back. Don't forget reset connection!
    ///
    /// ```rust
    /// use ntex_redis::{cmd, RedisConnector};
    ///
    /// #[ntex::main]
    /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
    ///     let redis = RedisConnector::new("127.0.0.1:6379").connect_simple().await?;
    ///    
    ///     let subscriber = redis.subscribe(cmd::Subscribe(vec!["test"]))?;
    ///     // do some work
    ///
    ///     // go back to normal client
    ///     let redis = subscriber.into_client();
    ///
    ///     // and reset connection, client may receive pending subscription messages instead of valid RESET response
    ///     if let Err(e) = redis.exec(cmd::Reset()).await {
    ///         println!("Error on reset connection: {}", e);      
    ///     };
    ///
    ///     Ok(())
    /// }
    /// ```
    pub fn into_client(self) -> SimpleClient {
        self.client
    }

    /// Send redis subscribe/unsubscribe command
    pub fn send<T: Command + PubSubCommand>(&self, cmd: T) -> Result<(), CommandError> {
        self.client.send(cmd)
    }

    /// 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)
    }
}