Add additional tracing instrumentation.

This commit lays the ground-work for making mini-redis a better
example of an instrumented tokio application. While it does not
go so far to turn mini-redis into a fully-featured guide for
using `tracing`, it ensures that people who use mini-redis as a
basis for experimenting with different subscribers get a more
complete experience out-of-the-box.

Author: jswrenn

Cargo.lock

@@ -47,6 +47,7 @@ impl Get {
     /// ```text
     /// GET key
     /// ```
+    #[instrument(level = "trace", name = "Get::parse_frames", skip(parse))]
     pub(crate) fn parse_frames(parse: &mut Parse) -> crate::Result<Get> {
         // The `GET` string has already been consumed. The next value is the
         // name of the key to get. If the next value is not a string or the
@@ -60,7 +61,14 @@ impl Get {
     ///
     /// The response is written to `dst`. This is called by the server in order
     /// to execute a received command.
-    #[instrument(skip(self, db, dst))]
+    #[instrument(
+        level = "trace",
+        name = "Get::apply",
+        skip(self, db, dst),
+        fields(
+            key = self.key.as_str(),
+        ),
+    )]
     pub(crate) async fn apply(self, db: &Db, dst: &mut Connection) -> crate::Result<()> {
         // Get the value from the shared database state
         let response = if let Some(value) = db.get(&self.key) {

src/cmd/get.rs

@@ -17,6 +17,7 @@ mod unknown;
 pub use unknown::Unknown;
 
 use crate::{Connection, Db, Frame, Parse, ParseError, Shutdown};
+use tracing::instrument;
 
 /// Enumeration of supported Redis commands.
 ///
@@ -41,6 +42,7 @@ impl Command {
     /// # Returns
     ///
     /// On success, the command value is returned, otherwise, `Err` is returned.
+    #[instrument(level = "trace", name = "Command::from_frame", skip(frame), err)]
     pub fn from_frame(frame: Frame) -> crate::Result<Command> {
         // The frame  value is decorated with `Parse`. `Parse` provides a
         // "cursor" like API which makes parsing the command easier.
@@ -87,6 +89,12 @@ impl Command {
     ///
     /// The response is written to `dst`. This is called by the server in order
     /// to execute a received command.
+    #[instrument(
+        level = "trace",
+        name = "Command::apply",
+        skip(self, db, dst, shutdown),
+        err
+    )]
     pub(crate) async fn apply(
         self,
         db: &Db,

src/cmd/mod.rs

@@ -39,6 +39,7 @@ impl Ping {
     /// ```text
     /// PING [message]
     /// ```
+    #[instrument(level = "trace", name = "Ping::parse_frames", skip(parse))]
     pub(crate) fn parse_frames(parse: &mut Parse) -> crate::Result<Ping> {
         match parse.next_string() {
             Ok(msg) => Ok(Ping::new(Some(msg))),
@@ -51,7 +52,13 @@ impl Ping {
     ///
     /// The response is written to `dst`. This is called by the server in order
     /// to execute a received command.
-    #[instrument(skip(self, dst))]
+    #[instrument(
+        name = "Ping::apply",
+        skip(self, dst),
+        fields(
+            ?msg = self.msg,
+        ),
+    )]
     pub(crate) async fn apply(self, dst: &mut Connection) -> crate::Result<()> {
         let response = match self.msg {
             None => Frame::Simple("PONG".to_string()),

src/cmd/ping.rs

@@ -1,6 +1,7 @@
 use crate::{Connection, Db, Frame, Parse};
 
 use bytes::Bytes;
+use tracing::instrument;
 
 /// Posts a message to the given channel.
 ///
@@ -47,6 +48,7 @@ impl Publish {
     /// ```text
     /// PUBLISH channel message
     /// ```
+    #[instrument(level = "trace", name = "Publish::parse_frames", skip(parse))]
     pub(crate) fn parse_frames(parse: &mut Parse) -> crate::Result<Publish> {
         // The `PUBLISH` string has already been consumed. Extract the `channel`
         // and `message` values from the frame.
@@ -64,6 +66,14 @@ impl Publish {
     ///
     /// The response is written to `dst`. This is called by the server in order
     /// to execute a received command.
+    #[instrument(
+        level = "trace",
+        name = "Publish::apply",
+        skip(self, db, dst),
+        fields(
+            channel = self.channel.as_str(),
+        ),
+    )]
     pub(crate) async fn apply(self, db: &Db, dst: &mut Connection) -> crate::Result<()> {
         // The shared state contains the `tokio::sync::broadcast::Sender` for
         // all active channels. Calling `db.publish` dispatches the message into

src/cmd/publish.rs

@@ -3,7 +3,7 @@ use crate::{Connection, Db, Frame};
 
 use bytes::Bytes;
 use std::time::Duration;
-use tracing::{debug, instrument};
+use tracing::instrument;
 
 /// Set `key` to hold the string `value`.
 ///
@@ -77,6 +77,7 @@ impl Set {
     /// ```text
     /// SET key value [EX seconds|PX milliseconds]
     /// ```
+    #[instrument(level = "trace", name = "Set::parse_frames", skip(parse))]
     pub(crate) fn parse_frames(parse: &mut Parse) -> crate::Result<Set> {
         use ParseError::EndOfStream;
 
@@ -124,14 +125,21 @@ impl Set {
     ///
     /// The response is written to `dst`. This is called by the server in order
     /// to execute a received command.
-    #[instrument(skip(self, db, dst))]
+    #[instrument(
+        level = "trace",
+        name = "Set::apply",
+        skip(self, db, dst),
+        fields(
+            key = self.key.as_str(),
+            ?expire = self.expire.as_ref().map(Duration::as_secs_f64),
+        ),
+    )]
     pub(crate) async fn apply(self, db: &Db, dst: &mut Connection) -> crate::Result<()> {
         // Set the value in the shared database state.
         db.set(self.key, self.value, self.expire);
 
         // Create a success response and write it to `dst`.
         let response = Frame::Simple("OK".to_string());
-        debug!(?response);
         dst.write_frame(&response).await?;
 
         Ok(())

src/cmd/set.rs

@@ -6,6 +6,7 @@ use std::pin::Pin;
 use tokio::select;
 use tokio::sync::broadcast;
 use tokio_stream::{Stream, StreamExt, StreamMap};
+use tracing::instrument;
 
 /// Subscribes the client to one or more channels.
 ///
@@ -60,6 +61,7 @@ impl Subscribe {
     /// ```text
     /// SUBSCRIBE channel [channel ...]
     /// ```
+    #[instrument(level = "trace", name = "Subscribe::parse_frames", skip(parse))]
     pub(crate) fn parse_frames(parse: &mut Parse) -> crate::Result<Subscribe> {
         use ParseError::EndOfStream;
 
@@ -99,6 +101,14 @@ impl Subscribe {
     /// are updated accordingly.
     ///
     /// [here]: https://redis.io/topics/pubsub
+    #[instrument(
+        level = "trace",
+        name = "Suscribe::apply",
+        skip(self, db, dst),
+        fields(
+            channels = "UNIMPLEMENTED", // FIXME
+        ),
+    )]
     pub(crate) async fn apply(
         mut self,
         db: &Db,

src/cmd/subscribe.rs

@@ -25,7 +25,14 @@ impl Unknown {
     /// Responds to the client, indicating the command is not recognized.
     ///
     /// This usually means the command is not yet implemented by `mini-redis`.
-    #[instrument(skip(self, dst))]
+    #[instrument(
+        level = "trace",
+        name = "Unknown::apply",
+        skip(self, dst),
+        fields(
+            command_name = self.command_name.as_str(),
+        ),
+    )]
     pub(crate) async fn apply(self, dst: &mut Connection) -> crate::Result<()> {
         let response = Frame::Error(format!("ERR unknown command '{}'", self.command_name));
 

src/cmd/unknown.rs

@@ -1,9 +1,11 @@
 use crate::frame::{self, Frame};
 
 use bytes::{Buf, BytesMut};
-use std::io::{self, Cursor};
+use std::io::{self, Cursor, ErrorKind::ConnectionReset};
+use std::net::SocketAddr;
 use tokio::io::{AsyncReadExt, AsyncWriteExt, BufWriter};
 use tokio::net::TcpStream;
+use tracing::warn;
 
 /// Send and receive `Frame` values from a remote peer.
 ///
@@ -28,6 +30,16 @@ pub struct Connection {
     buffer: BytesMut,
 }
 
+/// The result of [`Connection::maybe_read_bytes`].
+enum ConnectionState {
+    /// The connection was gracefully closed when reading was attempted.
+    Closed,
+    /// The connection was open when reading was attempted.
+    Open,
+    /// The connection was abruptly reset by the peer when reading was attempted.
+    Reset,
+}
+
 impl Connection {
     /// Create a new `Connection`, backed by `socket`. Read and write buffers
     /// are initialized.
@@ -42,6 +54,11 @@ impl Connection {
         }
     }
 
+    /// Returns the remote address that this connection is bound to.
+    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
+        self.stream.get_ref().peer_addr()
+    }
+
     /// Read a single `Frame` value from the underlying stream.
     ///
     /// The function waits until it has retrieved enough data to parse a frame.
@@ -63,23 +80,40 @@ impl Connection {
 
             // There is not enough buffered data to read a frame. Attempt to
             // read more data from the socket.
-            //
-            // On success, the number of bytes is returned. `0` indicates "end
-            // of stream".
-            if 0 == self.stream.read_buf(&mut self.buffer).await? {
-                // The remote closed the connection. For this to be a clean
-                // shutdown, there should be no data in the read buffer. If
-                // there is, this means that the peer closed the socket while
-                // sending a frame.
-                if self.buffer.is_empty() {
+            match self.maybe_read_bytes().await? {
+                ConnectionState::Open => continue,
+                ConnectionState::Closed | ConnectionState::Reset => {
+                    if !self.buffer.is_empty() {
+                        warn! {
+                            incomplete =? self.buffer,
+                            "connection closed with incomplete frame"
+                        };
+                    }
                     return Ok(None);
-                } else {
-                    return Err("connection reset by peer".into());
                 }
             }
         }
     }
 
+    /// Attempt to read bytes from the connection.
+    async fn maybe_read_bytes(&mut self) -> io::Result<ConnectionState> {
+        match self.stream.read_buf(&mut self.buffer).await {
+            // the connection was closed gracefully
+            Ok(0) => Ok(ConnectionState::Closed),
+            // the connection is still open
+            Ok(_) => Ok(ConnectionState::Open),
+            // the connection was closed abruptly by the peer
+            Err(e) if e.kind() == ConnectionReset => {
+                warn! {
+                    "connection closed abruptly by peer"
+                };
+                Ok(ConnectionState::Reset)
+            }
+            // reading failed for some other reason
+            Err(err) => Err(err),
+        }
+    }
+
     /// Tries to parse a frame from the buffer. If the buffer contains enough
     /// data, the frame is returned and the data removed from the buffer. If not
     /// enough data has been buffered yet, `Ok(None)` is returned. If the

src/connection.rs

@@ -4,7 +4,7 @@ use tokio::time::{self, Duration, Instant};
 use bytes::Bytes;
 use std::collections::{BTreeMap, HashMap};
 use std::sync::{Arc, Mutex};
-use tracing::debug;
+use tracing::{debug, instrument};
 
 /// A wrapper around a `Db` instance. This exists to allow orderly cleanup
 /// of the `Db` by signalling the background purge task to shut down when
@@ -150,6 +150,7 @@ impl Db {
     /// Returns `None` if there is no value associated with the key. This may be
     /// due to never having assigned a value to the key or a previously assigned
     /// value expired.
+    #[instrument(level = "trace", name = "Db::get", skip(self))]
     pub(crate) fn get(&self, key: &str) -> Option<Bytes> {
         // Acquire the lock, get the entry and clone the value.
         //
@@ -163,6 +164,7 @@ impl Db {
     /// Duration.
     ///
     /// If a value is already associated with the key, it is removed.
+    #[instrument(level = "trace", name = "Db::set", skip(self, value))]
     pub(crate) fn set(&self, key: String, value: Bytes, expire: Option<Duration>) {
         let mut state = self.shared.state.lock().unwrap();
 
@@ -231,6 +233,7 @@ impl Db {
     ///
     /// The returned `Receiver` is used to receive values broadcast by `PUBLISH`
     /// commands.
+    #[instrument(level = "trace", name = "Db::subscribe", skip(self, key))]
     pub(crate) fn subscribe(&self, key: String) -> broadcast::Receiver<Bytes> {
         use std::collections::hash_map::Entry;
 
@@ -262,6 +265,7 @@ impl Db {
 
     /// Publish a message to the channel. Returns the number of subscribers
     /// listening on the channel.
+    #[instrument(level = "trace", name = "Db::publish", skip(self, value))]
     pub(crate) fn publish(&self, key: &str, value: Bytes) -> usize {
         let state = self.shared.state.lock().unwrap();
 
@@ -279,6 +283,7 @@ impl Db {
 
     /// Signals the purge background task to shut down. This is called by the
     /// `DbShutdown`s `Drop` implementation.
+    #[instrument(name = "Db::shutdown_purge_task", skip(self))]
     fn shutdown_purge_task(&self) {
         // The background task must be signaled to shut down. This is done by
         // setting `State::shutdown` to `true` and signalling the task.
@@ -296,6 +301,7 @@ impl Db {
 impl Shared {
     /// Purge all expired keys and return the `Instant` at which the **next**
     /// key will expire. The background task will sleep until this instant.
+    #[instrument(name = "Shared::purge_expired_keys", skip(self))]
     fn purge_expired_keys(&self) -> Option<Instant> {
         let mut state = self.state.lock().unwrap();
 
@@ -323,6 +329,10 @@ impl Shared {
             }
 
             // The key expired, remove it
+            debug! {
+                key = &key.as_str(),
+                "purged_expired_key",
+            }
             state.entries.remove(key);
             state.expirations.remove(&(when, id));
         }
@@ -352,6 +362,7 @@ impl State {
 ///
 /// Wait to be notified. On notification, purge any expired keys from the shared
 /// state handle. If `shutdown` is set, terminate the task.
+#[instrument(skip(shared))]
 async fn purge_expired_tasks(shared: Arc<Shared>) {
     // If the shutdown flag is set, then the task should exit.
     while !shared.is_shutdown() {