Add Support for Tracing

**Context**
The [`tracing`](https://docs.rs/tracing/latest/tracing/index.html) crate enables more robust instrumentation
of Rust programs compared to basic logging. With this patch, we introduce the `tracing` crate
into AKD and only enable it when the `tracing` feature is specified. By default,
the feature is disabled and we continue to leverage basic logging.

In addition to adding `tracing` based logging and instrumentation throughout the
`akd` lib (i.e. not `akd_core`), various grammatical and organizational improvements
were made in areas where tracing was being added. Notably, the `log_metrics` functions
which exist throughout the storage layer were updated to simply log with `info` level
instead of taking an argument to specify the level. Rationale being that the `log_metrics`
functions are only called when the `runtime_metrics` feature is enabled and `info` is
a fair median to assume.

**Testing**
Since no major functional changes were made, the changes were tested via existing
automated tests with various different feature flags being passed to toggle `tracing`
on and off.

Author: Dillon George

Diff for: akd/Cargo.toml

@@ -14,8 +14,26 @@ readme = "../README.md"
 whatsapp_v1 = ["akd_core/whatsapp_v1"]
 experimental = ["akd_core/experimental"]
 
-rand = ["dep:rand"]
+# Default features mix (experimental + audit-proof protobuf mgmt support)
+default = [
+    "public_auditing",
+    "parallel_vrf",
+    "parallel_insert",
+    "preload_history",
+    "greedy_lookup_preload",
+    "experimental",
+]
+
 bench = ["experimental", "public_tests", "tokio/rt-multi-thread"]
+# Greedy loading of lookup proof nodes
+greedy_lookup_preload = []
+public_auditing = ["dep:protobuf", "akd_core/protobuf"]
+# Parallelize node insertion during publish
+parallel_insert = []
+# Parallelize VRF calculations during publish
+parallel_vrf = ["akd_core/parallel_vrf"]
+# Enable pre-loading of the nodes when generating history proofs
+preload_history = []
 public_tests = [
     "rand",
     "dep:colored",
@@ -25,30 +43,14 @@ public_tests = [
     "akd_core/rand",
     "dep:paste",
 ]
-public_auditing = ["dep:protobuf", "akd_core/protobuf"]
-serde_serialization = ["dep:serde", "akd_core/serde_serialization"]
+rand = ["dep:rand"]
 # Collect runtime metrics on db access calls + timing
 runtime_metrics = []
-# Parallelize VRF calculations during publish
-parallel_vrf = ["akd_core/parallel_vrf"]
-# Parallelize node insertion during publish
-parallel_insert = []
-# Enable pre-loading of the nodes when generating history proofs
-preload_history = []
+serde_serialization = ["dep:serde", "akd_core/serde_serialization"]
 # TESTING ONLY: Artifically slow the in-memory database (for benchmarking)
 slow_internal_db = []
-# Greedy loading of lookup proof nodes
-greedy_lookup_preload = []
-
-# Default features mix (experimental + audit-proof protobuf mgmt support)
-default = [
-    "public_auditing",
-    "parallel_vrf",
-    "parallel_insert",
-    "preload_history",
-    "greedy_lookup_preload",
-    "experimental",
-]
+# Tracing instrumentation
+tracing = ["dep:tracing"]
 
 [dependencies]
 ## Required dependencies ##
@@ -63,12 +65,13 @@ log = { version = "0.4", features = ["kv_unstable"] }
 tokio = { version = "1", features = ["sync", "time", "rt"] }
 
 ## Optional dependencies ##
-serde = { version = "1", features = ["derive"], optional = true }
-rand = { version = "0.8", optional = true }
 colored = { version = "2", optional = true }
 once_cell = { version = "1", optional = true }
-protobuf = { version = "3", optional = true }
 paste = { version = "1", optional = true }
+protobuf = { version = "3", optional = true }
+rand = { version = "0.8", optional = true }
+serde = { version = "1", features = ["derive"], optional = true }
+tracing = {version = "0.1.40", optional = true }
 
 [dev-dependencies]
 criterion = "0.5"

Diff for: akd/src/append_only_zks.rs

@@ -9,6 +9,7 @@
 
 use crate::hash::EMPTY_DIGEST;
 use crate::helper_structs::LookupInfo;
+use crate::log_shim::info;
 use crate::storage::manager::StorageManager;
 use crate::storage::types::StorageType;
 use crate::tree_node::{
@@ -22,15 +23,17 @@ use crate::{
     AppendOnlyProof, AzksElement, AzksValue, Digest, Direction, MembershipProof, NodeLabel,
     NonMembershipProof, PrefixOrdering, SiblingProof, SingleAppendOnlyProof, SizeOf, ARITY,
 };
+
+
 use async_recursion::async_recursion;
-use log::info;
 use std::cmp::Ordering;
 #[cfg(feature = "greedy_lookup_preload")]
 use std::collections::HashSet;
 use std::convert::TryFrom;
 use std::marker::Sync;
 use std::ops::Deref;
 
+
 /// The default azks key
 pub const DEFAULT_AZKS_KEY: u8 = 1u8;
 
@@ -802,7 +805,7 @@ impl Azks {
                     load_count
                 );
             }
-            storage.log_metrics(log::Level::Info).await;
+            storage.log_metrics().await;
 
             let (unchanged, leaves) = Self::get_append_only_proof_helper::<TC, _>(
                 latest_epoch,

Diff for: akd/src/auditor.rs

@@ -18,6 +18,7 @@ use crate::{
 };
 
 /// Verifies an audit proof, given start and end hashes for a merkle patricia tree.
+#[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
 pub async fn audit_verify<TC: Configuration>(
     hashes: Vec<Digest>,
     proof: AppendOnlyProof,
@@ -52,7 +53,8 @@ pub async fn audit_verify<TC: Configuration>(
     Ok(())
 }
 
-/// Helper for audit, verifies an append-only proof
+/// Helper for audit, verifies an append-only proof.
+#[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
 pub async fn verify_consecutive_append_only<TC: Configuration>(
     proof: &SingleAppendOnlyProof,
     start_hash: Digest,

Diff for: akd/src/directory.rs

@@ -11,6 +11,7 @@ use crate::append_only_zks::{Azks, InsertMode};
 use crate::ecvrf::{VRFKeyStorage, VRFPublicKey};
 use crate::errors::{AkdError, DirectoryError, StorageError};
 use crate::helper_structs::LookupInfo;
+use crate::log_shim::{error, info};
 use crate::storage::manager::StorageManager;
 use crate::storage::types::{DbRecord, ValueState, ValueStateRetrievalFlag};
 use crate::storage::Database;
@@ -23,11 +24,12 @@ use crate::VersionFreshness;
 use akd_core::configuration::Configuration;
 use akd_core::utils::get_marker_versions;
 use akd_core::verify::history::HistoryParams;
-use log::{error, info};
 use std::collections::{HashMap, HashSet};
 use std::marker::PhantomData;
 use std::sync::Arc;
 use tokio::sync::RwLock;
+#[cfg(feature = "tracing")]
+use tracing::Instrument;
 
 /// The representation of a auditable key directory
 pub struct Directory<TC, S: Database, V> {
@@ -64,6 +66,7 @@ where
     /// Creates a new (stateless) instance of a auditable key directory.
     /// Takes as input a pointer to the storage being used for this instance.
     /// The state is stored in the storage.
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn new(storage: StorageManager<S>, vrf: V) -> Result<Self, AkdError> {
         let azks = Directory::<TC, S, V>::get_azks_from_storage(&storage, false).await;
 
@@ -90,8 +93,9 @@ where
     ///
     /// Note that the vector of label-value pairs should not contain any entries with duplicate labels. This
     /// condition is explicitly checked, and an error will be returned if this is the case.
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all, fields(num_updates = updates.len())))]
     pub async fn publish(&self, updates: Vec<(AkdLabel, AkdValue)>) -> Result<EpochHash, AkdError> {
-        // The guard will be dropped at the end of the publish
+        // The guard will be dropped at the end of the publish operation
         let _guard = self.cache_lock.read().await;
 
         // Check for duplicate labels and return an error if any are encountered
@@ -254,6 +258,7 @@ where
     ///
     /// Returns [Ok((LookupProof, EpochHash))] upon successful generation for the latest version
     /// of the target label's state. [Err(_)] otherwise
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn lookup(&self, akd_label: AkdLabel) -> Result<(LookupProof, EpochHash), AkdError> {
         // The guard will be dropped at the end of the proof generation
         let _guard = self.cache_lock.read().await;
@@ -281,6 +286,7 @@ where
     ///   from bulk lookup proof generation, as it has its own preloading operation
     ///
     /// Returns [Ok(LookupProof)] if the proof generation succeeded, [Err(_)] otherwise
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     async fn lookup_with_info(
         &self,
         current_azks: &Azks,
@@ -351,6 +357,7 @@ where
 
     // TODO(eoz): Call proof generations async
     /// Allows efficient batch lookups by preloading necessary nodes for the lookups.
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn batch_lookup(
         &self,
         akd_labels: &[AkdLabel],
@@ -392,6 +399,7 @@ where
         Ok((lookup_proofs, root_hash))
     }
 
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     async fn build_lookup_info(&self, latest_st: &ValueState) -> Result<LookupInfo, AkdError> {
         let akd_label = &latest_st.username;
         // Need to account for the case where the latest state is
@@ -419,6 +427,7 @@ where
         })
     }
 
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     async fn get_lookup_info(
         &self,
         akd_label: AkdLabel,
@@ -449,13 +458,21 @@ where
     /// this function returns all the values ever associated with it,
     /// and the epoch at which each value was first committed to the server state.
     /// It also returns the proof of the latest version being served at all times.
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn key_history(
         &self,
         akd_label: &AkdLabel,
         params: HistoryParams,
     ) -> Result<(HistoryProof, EpochHash), AkdError> {
         // The guard will be dropped at the end of the proof generation
+        #[cfg(not(feature = "tracing"))]
         let _guard = self.cache_lock.read().await;
+        #[cfg(feature = "tracing")]
+        let _guard = self
+            .cache_lock
+            .read()
+            .instrument(tracing::trace_span!("cache_lock.read"))
+            .await;
 
         let current_azks = self.retrieve_azks().await?;
         let current_epoch = current_azks.get_latest_epoch();
@@ -616,9 +633,24 @@ where
                 {
                     // acquire a singleton lock prior to flushing the cache to assert that no
                     // cache accesses are underway (i.e. publish/proof generations/etc)
+                    #[cfg(not(feature = "tracing"))]
                     let _guard = self.cache_lock.write().await;
+                    #[cfg(feature = "tracing")]
+                    let _guard = self
+                        .cache_lock
+                        .write()
+                        .instrument(tracing::trace_span!("cache_lock.write"))
+                        .await;
+
                     // flush the cache in its entirety
+                    #[cfg(not(feature = "tracing"))]
                     self.storage.flush_cache().await;
+                    #[cfg(feature = "tracing")]
+                    self.storage
+                        .flush_cache()
+                        .instrument(tracing::trace_span!("flush_cache"))
+                        .await;
+
                     // re-fetch the azks to load it into cache so when we release the cache lock
                     // others will see the new AZKS loaded up and ready
                     last =
@@ -643,13 +675,21 @@ where
 
     /// Returns an [AppendOnlyProof] for the leaves inserted into the underlying tree between
     /// the epochs `audit_start_ep` and `audit_end_ep`.
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all, fields(start_epoch = audit_start_ep, end_epoch = audit_end_ep)))]
     pub async fn audit(
         &self,
         audit_start_ep: u64,
         audit_end_ep: u64,
     ) -> Result<AppendOnlyProof, AkdError> {
         // The guard will be dropped at the end of the proof generation
+        #[cfg(not(feature = "tracing"))]
         let _guard = self.cache_lock.read().await;
+        #[cfg(feature = "tracing")]
+        let _guard = self
+            .cache_lock
+            .read()
+            .instrument(tracing::trace_span!("cache_lock.read"))
+            .await;
 
         let current_azks = self.retrieve_azks().await?;
         let current_epoch = current_azks.get_latest_epoch();
@@ -673,10 +713,12 @@ where
     }
 
     /// Retrieves the [Azks]
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub(crate) async fn retrieve_azks(&self) -> Result<Azks, crate::errors::AkdError> {
         Directory::<TC, S, V>::get_azks_from_storage(&self.storage, false).await
     }
 
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all, fields(ignore_cache = ignore_cache)))]
     async fn get_azks_from_storage(
         storage: &StorageManager<S>,
         ignore_cache: bool,
@@ -704,10 +746,12 @@ where
     /// HELPERS ///
 
     /// Use this function to retrieve the [VRFPublicKey] for this AKD.
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn get_public_key(&self) -> Result<VRFPublicKey, AkdError> {
         Ok(self.vrf.get_vrf_public_key().await?)
     }
 
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     async fn create_single_update_proof(
         &self,
         akd_label: &AkdLabel,
@@ -770,6 +814,7 @@ where
     }
 
     /// Gets the root hash at the current epoch.
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn get_epoch_hash(&self) -> Result<EpochHash, AkdError> {
         let current_azks = self.retrieve_azks().await?;
         let latest_epoch = current_azks.get_latest_epoch();
@@ -823,11 +868,13 @@ where
     }
 
     /// Read-only access to [Directory::lookup](Directory::lookup).
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn lookup(&self, uname: AkdLabel) -> Result<(LookupProof, EpochHash), AkdError> {
         self.0.lookup(uname).await
     }
 
     /// Read-only access to [Directory::batch_lookup](Directory::batch_lookup).
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn batch_lookup(
         &self,
         unames: &[AkdLabel],
@@ -836,6 +883,7 @@ where
     }
 
     /// Read-only access to [Directory::key_history](Directory::key_history).
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn key_history(
         &self,
         uname: &AkdLabel,
@@ -845,6 +893,7 @@ where
     }
 
     /// Read-only access to [Directory::poll_for_azks_changes](Directory::poll_for_azks_changes).
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn poll_for_azks_changes(
         &self,
         period: tokio::time::Duration,
@@ -854,6 +903,7 @@ where
     }
 
     /// Read-only access to [Directory::audit](Directory::audit).
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn audit(
         &self,
         audit_start_ep: u64,
@@ -863,11 +913,13 @@ where
     }
 
     /// Read-only access to [Directory::get_epoch_hash].
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn get_epoch_hash(&self) -> Result<EpochHash, AkdError> {
         self.0.get_epoch_hash().await
     }
 
     /// Read-only access to [Directory::get_public_key](Directory::get_public_key).
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all))]
     pub async fn get_public_key(&self) -> Result<VRFPublicKey, AkdError> {
         self.0.get_public_key().await
     }

Diff for: akd/src/lib.rs

@@ -254,7 +254,7 @@
 //! First, it encodes a [`HistoryParams`]. Note that the same argument for [`HistoryParams`] that was used to generate
 //! the key history proof must also be used to verify the proof. Otherwise, verification may fail.
 //!
-//! [`HistoryVerificationParams`] also allows the consumer to specify whether or not a "tombstoned" value should be
+//! [`HistoryVerificationParams`] also allows the consumer to specify whether a "tombstoned" value should be
 //! accepted in place of a valid value for the corresponding entry. This is useful
 //! in scenarios where the consumer wishes to verify that a particular entry exists,
 //! but does not care about the value associated with it. The default behavior is to
@@ -499,6 +499,19 @@ pub mod helper_structs;
 pub mod storage;
 pub mod tree_node;
 
+/// Shim module to group logging-related macros more easily
+/// when switching between [log] and [tracing].
+pub mod log_shim {
+    #[cfg(feature = "tracing")]
+    pub use {
+        tracing::{debug, error, info, trace, warn},
+    };
+    #[cfg(not(feature = "tracing"))]
+    pub use {
+        log::{debug, error, info, trace, warn},
+    };
+}
+
 #[cfg(feature = "public_auditing")]
 pub mod local_auditing;