nitox 0.1.9

Nitox is a `tokio`-based client for NATS.
Documentation 
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

use futures::{
    future::{self, Either},
    prelude::*,
};
use parking_lot::RwLock;
use std::{net::SocketAddr, sync::Arc};
use tokio_executor;

use error::NatsError;
use protocol::Op;

use super::connection_inner::NatsConnectionInner;

macro_rules! reco {
    ($conn:ident) => {
        *$conn.state.write() = NatsConnectionState::Disconnected;

        tokio_executor::spawn($conn.reconnect().map_err(|e| {
            debug!(target: "nitox", "Reconnection error: {}", e);
            ()
        }));
    };
}

/// State of the raw connection
#[derive(Debug, Clone, Copy, PartialEq)]
pub(crate) enum NatsConnectionState {
    Connected,
    Reconnecting,
    Disconnected,
}

/// Represents a connection to a NATS server. Implements `Sink` and `Stream`
#[derive(Debug)]
pub struct NatsConnection {
    /// indicates if the connection is made over TLS
    pub(crate) is_tls: bool,
    /// Server standardized IP address
    pub(crate) addr: SocketAddr,
    /// Host of the server; Only used if connecting to a TLS-enabled server
    pub(crate) host: Option<String>,
    /// Inner dual `Stream`/`Sink` of the TCP connection
    pub(crate) inner: Arc<RwLock<NatsConnectionInner>>,
    /// Current state of the connection
    pub(crate) state: Arc<RwLock<NatsConnectionState>>,
}

impl NatsConnection {
    /// Tries to reconnect once to the server; Only used internally. Blocks polling during reconnecting
    /// by forcing the object to return `Async::NotReady`/`AsyncSink::NotReady`
    fn reconnect(&self) -> impl Future<Item = (), Error = NatsError> {
        *self.state.write() = NatsConnectionState::Reconnecting;

        let inner_arc = Arc::clone(&self.inner);
        let inner_state = Arc::clone(&self.state);
        let is_tls = self.is_tls;
        let maybe_host = self.host.clone();
        NatsConnectionInner::connect_tcp(&self.addr)
            .and_then(move |socket| {
                if is_tls {
                    Either::A(
                        // This unwrap is safe because the value would always be present if `is_tls` is true
                        NatsConnectionInner::upgrade_tcp_to_tls(&maybe_host.unwrap(), socket)
                            .map(NatsConnectionInner::from),
                    )
                } else {
                    Either::B(future::ok(NatsConnectionInner::from(socket)))
                }
            }).and_then(move |inner| {
                {
                    *inner_arc.write() = inner;
                    *inner_state.write() = NatsConnectionState::Connected;
                }
                debug!(target: "nitox", "Successfully swapped reconnected underlying connection");
                Ok(())
            })
    }
}

impl Sink for NatsConnection {
    type SinkError = NatsError;
    type SinkItem = Op;

    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        if match self.state.try_read() {
            Some(state) => *state != NatsConnectionState::Connected,
            _ => true,
        } {
            return Ok(AsyncSink::NotReady(item));
        }

        if let Some(mut inner) = self.inner.try_write() {
            match inner.start_send(item.clone()) {
                Err(NatsError::ServerDisconnected(_)) => {
                    reco!(self);
                    Ok(AsyncSink::NotReady(item))
                }
                poll_res => poll_res,
            }
        } else {
            Ok(AsyncSink::NotReady(item))
        }
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        if match self.state.try_read() {
            Some(state) => *state != NatsConnectionState::Connected,
            _ => true,
        } {
            return Ok(Async::NotReady);
        }

        if let Some(mut inner) = self.inner.try_write() {
            match inner.poll_complete() {
                Err(NatsError::ServerDisconnected(_)) => {
                    reco!(self);
                    Ok(Async::NotReady)
                }
                poll_res => poll_res,
            }
        } else {
            Ok(Async::NotReady)
        }
    }
}

impl Stream for NatsConnection {
    type Error = NatsError;
    type Item = Op;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        if match self.state.try_read() {
            Some(state) => *state != NatsConnectionState::Connected,
            _ => true,
        } {
            return Ok(Async::NotReady);
        }

        if let Some(mut inner) = self.inner.try_write() {
            match inner.poll() {
                Err(NatsError::ServerDisconnected(_)) => {
                    reco!(self);
                    Ok(Async::NotReady)
                }
                poll_res => poll_res,
            }
        } else {
            Ok(Async::NotReady)
        }
    }
}