Async client (unchanged API)

Cargo.toml

@@ -39,6 +39,9 @@ log = "0.4.8"
 nkeys = "0.0.9"
 base64-url = "1.1.13"
 once_cell = "1.3.1"
+smol = "0.1.4"
+piper = "0.1.1"
+futures = "0.3.4"
 
 [dev-dependencies]
 criterion = "0.3.2"

examples/new-client.rs

@@ -0,0 +1,24 @@
+use std::io;
+use std::thread;
+use std::time::Duration;
+
+use nats::new_client::Connection;
+
+fn main() -> io::Result<()> {
+    // Useful commands for testing:
+    // nats-sub -s nats://demo.nats.io hello
+    // nats-pub -s nats://demo.nats.io hello 'hi from nats-pub'
+
+    let mut nc = Connection::connect("demo.nats.io:4222")?;
+    let mut sub = nc.subscribe("hello");
+
+    thread::sleep(Duration::from_secs(1));
+
+    nc.publish("hello", "hi from new-client")?;
+    nc.flush()?;
+
+    loop {
+        let msg = sub.next_msg()?;
+        println!("{}", msg);
+    }
+}

src/lib.rs

@@ -9,8 +9,9 @@
 //!
 //! [https://nats.io/]: https://nats.io/
 
-#![cfg_attr(test, deny(warnings))]
-#![deny(
+#![recursion_limit = "1024"]
+// #![cfg_attr(test, deny(warnings))]
+#![warn(
     missing_docs,
     future_incompatible,
     nonstandard_style,
@@ -92,6 +93,8 @@ mod parser;
 mod secure_wipe;
 mod shared_state;
 
+pub mod new_client;
+
 #[cfg(feature = "fault_injection")]
 mod fault_injection;
 

src/new_client/client.rs

@@ -0,0 +1,223 @@
+use std::collections::VecDeque;
+use std::io::{self, Error, ErrorKind};
+use std::net::TcpStream;
+use std::thread;
+use std::time::{Duration, Instant};
+
+use futures::{
+    channel::{mpsc, oneshot},
+    io::{BufReader, BufWriter},
+    prelude::*,
+};
+use piper::Arc;
+use smol::{Async, Timer};
+
+use crate::connect::ConnectInfo;
+use crate::new_client::decoder::{decode, ServerOp};
+use crate::new_client::encoder::{encode, ClientOp};
+use crate::Message;
+
+/// An operation enqueued by the user of this crate.
+///
+/// These operations are enqueued by `Connection` or `Subscription` and are then handled by the
+/// client thread.
+pub(crate) enum UserOp {
+    /// Publish a message.
+    Pub {
+        subject: String,
+        reply_to: Option<String>,
+        payload: Vec<u8>,
+    },
+
+    /// Subscribe to a subject.
+    Sub {
+        subject: String,
+        queue_group: Option<String>,
+        sid: usize,
+        messages: mpsc::UnboundedSender<Message>,
+    },
+
+    /// Unsubscribe from a subject.
+    Unsub { sid: usize, max_msgs: Option<u64> },
+
+    /// Send a ping and wait for a pong.
+    Ping { pong: oneshot::Sender<()> },
+
+    /// Close the connection.
+    Close,
+}
+
+/// Spawns a client thread.
+pub(crate) fn spawn(
+    url: &str,
+    user_ops: mpsc::UnboundedReceiver<UserOp>,
+) -> thread::JoinHandle<io::Result<()>> {
+    let url = url.to_string();
+    thread::spawn(move || smol::run(client(&url, user_ops)))
+}
+
+/// Runs the main loop for a client.
+async fn client(url: &str, mut user_ops: mpsc::UnboundedReceiver<UserOp>) -> io::Result<()> {
+    let stream = Arc::new(Async::<TcpStream>::connect(url).await?);
+
+    // TODO(stjepang): Make this option configurable.
+    let flush_timeout = Duration::from_millis(100);
+
+    // Bytes written to the server are buffered and periodically flushed.
+    let mut next_flush = Instant::now() + flush_timeout;
+    let mut writer = BufWriter::new(stream.clone());
+
+    // Create an endless stream parsing operations from the server.
+    let mut server_ops = stream::try_unfold(BufReader::new(stream), |mut stream| async {
+        // Decode a single operation.
+        let op = decode(&mut stream).await?;
+        io::Result::Ok(Some((op, stream)))
+    })
+    .boxed();
+
+    // Expect an INFO message.
+    let mut server_info = match server_ops
+        .try_next()
+        .await?
+        .expect("end of what should be an infinite stream")
+    {
+        ServerOp::Info(server_info) => server_info,
+        _ => return Err(Error::new(ErrorKind::Other, "expected an INFO message")),
+    };
+
+    // Current subscriptions in the form `(subject, sid, messages)`.
+    let mut subscriptions: Vec<(String, usize, mpsc::UnboundedSender<Message>)> = Vec::new();
+
+    // Expected pongs and their notification channels.
+    let mut pongs: VecDeque<oneshot::Sender<()>> = VecDeque::new();
+
+    // Send a CONNECT operation to the server.
+    encode(
+        &mut writer,
+        ClientOp::Connect(ConnectInfo {
+            tls_required: false,
+            name: None,
+            pedantic: false,
+            verbose: false,
+            lang: crate::LANG.to_string(),
+            version: crate::VERSION.to_string(),
+            user: None,
+            pass: None,
+            auth_token: None,
+            user_jwt: None,
+            signature: None,
+            echo: true,
+        }),
+    )
+    .await?;
+
+    // Handle events in a loop.
+    loop {
+        futures::select! {
+            // An operation was received from the server.
+            res = server_ops.try_next().fuse() => {
+                let op = res?.expect("end of what should be an infinite stream");
+
+                match op {
+                    ServerOp::Info(new_server_info) => {
+                        server_info = new_server_info;
+                    }
+
+                    ServerOp::Ping => {
+                        // Send a PONG operation to the server.
+                        encode(&mut writer, ClientOp::Pong).await?;
+                    }
+
+                    ServerOp::Pong => {
+                        // Take the next expected pong from the queue.
+                        let pong = pongs.pop_front().expect("unexpected pong");
+
+                        // Complete the pong by sending a message into the channel.
+                        let _ = pong.send(());
+                    }
+
+                    ServerOp::Msg { subject, sid, reply_to, payload } => {
+                        // Send the message to matching subscriptions.
+                        for (_, _, messages) in
+                            subscriptions.iter().filter(|(_, s, _)| *s == sid)
+                        {
+                            let _ = messages.unbounded_send(Message {
+                                subject: subject.clone(),
+                                reply: reply_to.clone(),
+                                data: payload.clone(),
+                                responder: None,
+                            });
+                        }
+                    }
+
+                    ServerOp::Err(msg) => {
+                        log::error!("received -ERR '{}'", msg);
+                    }
+
+                    ServerOp::Unknown(line) => {
+                        log::warn!("unknown op: {}", line);
+                    }
+                }
+            }
+
+            // The user has enqueued an operation.
+            msg = user_ops.next().fuse() => {
+                match msg.expect("user_ops disconnected") {
+                    UserOp::Pub { subject, reply_to, payload } => {
+                        // Send a PUB operation to the server.
+                        encode(&mut writer, ClientOp::Pub {
+                            subject,
+                            reply_to: None,
+                            payload
+                        })
+                        .await?;
+                    }
+
+                    UserOp::Sub { subject, queue_group, sid, messages } => {
+                        // Add the subscription to the list.
+                        subscriptions.push((subject.clone(), sid, messages));
+
+                        // Send a SUB operation to the server.
+                        encode(&mut writer, ClientOp::Sub {
+                            subject,
+                            queue_group,
+                            sid,
+                        })
+                        .await?;
+                    }
+
+                    UserOp::Unsub { sid, max_msgs } => {
+                        // Remove the subscription from the list.
+                        subscriptions.retain(|(_, s, _)| *s != sid);
+
+                        // Send an UNSUB operation to the server.
+                        encode(&mut writer, ClientOp::Unsub {
+                            sid,
+                            max_msgs,
+                        })
+                        .await?;
+                    }
+
+                    UserOp::Ping { pong } => {
+                        // Send a PING operation to the server.
+                        encode(&mut writer, ClientOp::Ping).await?;
+
+                        // Record that we're expecting a pong.
+                        pongs.push_back(pong);
+                    }
+
+                    UserOp::Close => {
+                        // TODO(stjepang): Perhaps we should flush before closing abruptly.
+                        return Ok(());
+                    }
+                }
+            }
+
+            // Periodically flush writes to the server.
+            _ = Timer::at(next_flush).fuse() => {
+                writer.flush().await?;
+                next_flush = Instant::now() + flush_timeout;
+            }
+        }
+    }
+}

src/new_client/connection.rs

@@ -0,0 +1,105 @@
+use std::io::{self, Error, ErrorKind};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::thread;
+
+use futures::channel::{mpsc, oneshot};
+use smol::block_on;
+
+use crate::new_client::client::{self, UserOp};
+use crate::new_client::subscription::Subscription;
+
+/// A NATS client connection.
+pub struct Connection {
+    /// Enqueues user operations.
+    user_ops: mpsc::UnboundedSender<UserOp>,
+
+    /// Subscription ID generator.
+    sid_gen: AtomicUsize,
+
+    /// Thread running the main future.
+    thread: Option<thread::JoinHandle<io::Result<()>>>,
+}
+
+impl Connection {
+    /// Connects a NATS client.
+    pub fn connect(url: &str) -> io::Result<Connection> {
+        // Spawn a client thread.
+        let (sender, receiver) = mpsc::unbounded();
+        let thread = client::spawn(url, receiver);
+
+        // Connection handle controlling the client thread.
+        let mut conn = Connection {
+            user_ops: sender,
+            sid_gen: AtomicUsize::new(1),
+            thread: Some(thread),
+        };
+
+        // Flush to send a ping and wait for the connection to establish.
+        conn.flush()?;
+
+        // All good! The connection is now ready.
+        Ok(conn)
+    }
+
+    /// Publishes a message.
+    pub fn publish(&mut self, subject: &str, msg: impl AsRef<[u8]>) -> io::Result<()> {
+        let subject = subject.to_string();
+        let payload = msg.as_ref().to_vec();
+        let reply_to = None;
+
+        // Enqueue a PUB operation.
+        self.user_ops
+            .unbounded_send(UserOp::Pub {
+                subject,
+                reply_to,
+                payload,
+            })
+            .map_err(|err| Error::new(ErrorKind::ConnectionReset, err))?;
+
+        Ok(())
+    }
+
+    /// Creates a new subscriber.
+    pub fn subscribe(&mut self, subject: &str) -> Subscription {
+        let sid = self.sid_gen.fetch_add(1, Ordering::SeqCst);
+        Subscription::new(subject, sid, self.user_ops.clone())
+    }
+
+    /// Flushes by performing a round trip to the server.
+    pub fn flush(&mut self) -> io::Result<()> {
+        let (sender, receiver) = oneshot::channel();
+
+        // Enqueue a PING operation.
+        self.user_ops
+            .unbounded_send(UserOp::Ping { pong: sender })
+            .map_err(|err| Error::new(ErrorKind::ConnectionReset, err))?;
+
+        // Wait until the PONG operation is received.
+        let _ = block_on(receiver);
+
+        Ok(())
+    }
+
+    /// Close the connection.
+    pub fn close(&mut self) -> io::Result<()> {
+        if let Some(thread) = self.thread.take() {
+            // Enqueue a close operation.
+            let _ = self.user_ops.unbounded_send(UserOp::Close);
+
+            // Wait for the client thread to stop.
+            thread
+                .join()
+                .expect("client thread has panicked")
+                .map_err(|err| Error::new(ErrorKind::ConnectionReset, err))?;
+        }
+
+        Ok(())
+    }
+}
+
+impl Drop for Connection {
+    fn drop(&mut self) {
+        // Close the connection in case it hasn't been already.
+        let _ = self.close();
+    }
+}