Remove the epochs vector from HistoryTreeNode (#113)

* Remove the epochs vector from the history tree node since the vast majority of operations only require the birth and current epochs. (Only audit & key history require the full list of epochs a node was mutated in) Resolves #112 * Add mysql limit to limit select as small as possible * Version bump to 0.3.6 Co-authored-by: Sean Lawlor <[email protected]>
facebook · Dec 14, 2021 · 1e2a0cb · 1e2a0cb
1 parent 6ca9e97
commit 1e2a0cb
Show file tree

Hide file tree

Showing 8 changed files with 247 additions and 96 deletions.
diff --git a/akd/Cargo.toml b/akd/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "akd"
-version = "0.3.5"
+version = "0.3.6"
 authors = ["Harjasleen Malvai <[email protected]>", "Kevin Lewi <[email protected]>", "Sean Lawlor <[email protected]>"]
 description = "An implementation of an auditable key directory"
 license = "MIT OR Apache-2.0"

diff --git a/akd/src/history_tree_node.rs b/akd/src/history_tree_node.rs
@@ -55,8 +55,10 @@ pub(crate) type HistoryInsertionNode = (Direction, HistoryChildState);
 pub struct HistoryTreeNode {
     /// The binary label for this node
     pub label: NodeLabel,
-    /// The epochs this node was updated
-    pub epochs: Vec<u64>,
+    /// The last epoch this node was updated in
+    pub last_epoch: u64,
+    /// The epoch that this node was birthed in
+    pub birth_epoch: u64,
     /// The label of this node's parent
     pub parent: NodeLabel, // The root node is marked its own parent.
     /// The type of node: leaf root or interior.
@@ -111,18 +113,20 @@ impl Clone for HistoryTreeNode {
     fn clone(&self) -> Self {
         Self {
             label: self.label,
-            epochs: self.epochs.clone(),
+            last_epoch: self.last_epoch,
+            birth_epoch: self.birth_epoch,
             parent: self.parent,
             node_type: self.node_type,
         }
     }
 }
 
 impl HistoryTreeNode {
-    fn new(label: NodeLabel, parent: NodeLabel, node_type: NodeType) -> Self {
+    fn new(label: NodeLabel, parent: NodeLabel, node_type: NodeType, birth_epoch: u64) -> Self {
         HistoryTreeNode {
             label,
-            epochs: vec![],
+            birth_epoch,
+            last_epoch: birth_epoch,
             parent, // Root node is its own parent
             node_type,
         }
@@ -245,8 +249,7 @@ impl HistoryTreeNode {
 
                 debug!("BEGIN create new node");
                 let mut new_node =
-                    HistoryTreeNode::new(lcs_label, parent.label, NodeType::Interior);
-                new_node.epochs.push(epoch);
+                    HistoryTreeNode::new(lcs_label, parent.label, NodeType::Interior, epoch);
                 new_node.write_to_storage(storage).await?;
                 set_state_map(
                     storage,
@@ -464,11 +467,11 @@ impl HistoryTreeNode {
             match self.get_latest_epoch() {
                 Ok(latest) => {
                     if latest != epoch {
-                        self.epochs.push(epoch);
+                        self.last_epoch = epoch;
                     }
                 }
                 Err(_) => {
-                    self.epochs.push(epoch);
+                    self.last_epoch = epoch;
                 }
             }
             self.write_to_storage(storage).await?;
@@ -572,7 +575,7 @@ impl HistoryTreeNode {
     }
 
     pub(crate) fn get_birth_epoch(&self) -> u64 {
-        self.epochs[0]
+        self.birth_epoch
     }
 
     // gets the direction of node, i.e. if it's a left
@@ -617,14 +620,31 @@ impl HistoryTreeNode {
                 if self.get_birth_epoch() > epoch {
                     Err(HistoryTreeNodeError::NoChildInTreeAtEpoch(epoch, dir))
                 } else {
-                    let mut chosen_ep = self.get_birth_epoch();
-                    for existing_ep in &self.epochs {
-                        if *existing_ep <= epoch && *existing_ep > chosen_ep {
-                            chosen_ep = *existing_ep;
+                    let chosen_ep = {
+                        if self.last_epoch <= epoch {
+                            // the "last" updated epoch is <= epoch, so it is
+                            // the last valid state at this epoch
+                            Some(self.last_epoch)
+                        } else if self.birth_epoch == epoch {
+                            // we're looking at the state at the birth epoch
+                            Some(self.birth_epoch)
+                        } else {
+                            // Indeterminate, we are somewhere above the
+                            // birth epoch but we're less than the "last" epoch.
+                            // db query is necessary
+                            None
                         }
+                    };
+
+                    if let Some(ep) = chosen_ep {
+                        self.get_child_at_existing_epoch::<_, H>(storage, ep, direction)
+                            .await
+                    } else {
+                        let target_ep = storage.get_epoch_lte_epoch(self.label, epoch).await?;
+                        // DB query for the state <= this epoch value
+                        self.get_child_at_existing_epoch::<_, H>(storage, target_ep, direction)
+                            .await
                     }
-                    self.get_child_at_existing_epoch::<_, H>(storage, chosen_ep, direction)
-                        .await
                 }
             }
         }
@@ -652,13 +672,28 @@ impl HistoryTreeNode {
         if self.get_birth_epoch() > epoch {
             Err(HistoryTreeNodeError::NodeDidNotExistAtEp(self.label, epoch))
         } else {
-            let mut chosen_ep = self.get_birth_epoch();
-            for existing_ep in &self.epochs {
-                if *existing_ep <= epoch {
-                    chosen_ep = *existing_ep;
+            let chosen_ep = {
+                if self.last_epoch <= epoch {
+                    // the "last" updated epoch is <= epoch, so it is
+                    // the last valid state at this epoch
+                    Some(self.last_epoch)
+                } else if self.birth_epoch == epoch {
+                    // we're looking at the state at the birth epoch
+                    Some(self.birth_epoch)
+                } else {
+                    // Indeterminate, we are somewhere above the
+                    // birth epoch but we're less than the "last" epoch.
+                    // db query is necessary
+                    None
                 }
+            };
+            if let Some(ep) = chosen_ep {
+                self.get_state_at_existing_epoch(storage, ep).await
+            } else {
+                let target_ep = storage.get_epoch_lte_epoch(self.label, epoch).await?;
+                // DB query for the state <= this epoch value
+                self.get_state_at_existing_epoch(storage, target_ep).await
             }
-            self.get_state_at_existing_epoch(storage, chosen_ep).await
         }
     }
 
@@ -675,12 +710,7 @@ impl HistoryTreeNode {
     /* Functions for compression-related operations */
 
     pub(crate) fn get_latest_epoch(&self) -> Result<u64, HistoryTreeNodeError> {
-        match self.epochs.len() {
-            0 => Err(HistoryTreeNodeError::NodeCreatedWithoutEpochs(
-                self.label.get_val(),
-            )),
-            n => Ok(self.epochs[n - 1]),
-        }
+        Ok(self.last_epoch)
     }
 
     /////// Helpers /////////
@@ -739,9 +769,10 @@ pub(crate) async fn get_empty_root<H: Hasher, S: Storage + Send + Sync>(
     storage: &S,
     ep: Option<u64>,
 ) -> Result<HistoryTreeNode, HistoryTreeNodeError> {
-    let mut node = HistoryTreeNode::new(NodeLabel::root(), NodeLabel::root(), NodeType::Root);
+    let mut node = HistoryTreeNode::new(NodeLabel::root(), NodeLabel::root(), NodeType::Root, 0u64);
     if let Some(epoch) = ep {
-        node.epochs.push(epoch);
+        node.birth_epoch = epoch;
+        node.last_epoch = epoch;
         let new_state: HistoryNodeState =
             HistoryNodeState::new::<H>(NodeStateKey(node.label, epoch));
         set_state_map(storage, new_state).await?;
@@ -759,7 +790,8 @@ pub(crate) async fn get_leaf_node<H: Hasher, S: Storage + Sync + Send>(
 ) -> Result<HistoryTreeNode, HistoryTreeNodeError> {
     let node = HistoryTreeNode {
         label,
-        epochs: vec![birth_epoch],
+        birth_epoch,
+        last_epoch: birth_epoch,
         parent,
         node_type: NodeType::Leaf,
     };
@@ -782,7 +814,8 @@ pub(crate) async fn get_leaf_node_without_hashing<H: Hasher, S: Storage + Sync +
 ) -> Result<HistoryTreeNode, HistoryTreeNodeError> {
     let node = HistoryTreeNode {
         label,
-        epochs: vec![birth_epoch],
+        birth_epoch,
+        last_epoch: birth_epoch,
         parent,
         node_type: NodeType::Leaf,
     };

diff --git a/akd/src/storage/memory.rs b/akd/src/storage/memory.rs
@@ -266,6 +266,43 @@ impl Storage for AsyncInMemoryDatabase {
         }
         Ok(map)
     }
+
+    async fn get_epoch_lte_epoch(
+        &self,
+        node_label: crate::node_state::NodeLabel,
+        epoch_in_question: u64,
+    ) -> Result<u64, StorageError> {
+        let ids = (0..=epoch_in_question)
+            .map(|epoch| crate::node_state::NodeStateKey(node_label, epoch))
+            .collect::<Vec<_>>();
+        let data = self
+            .batch_get::<crate::node_state::HistoryNodeState>(ids)
+            .await?;
+        let mut epochs = data
+            .into_iter()
+            .map(|item| {
+                if let DbRecord::HistoryNodeState(state) = item {
+                    state.key.1
+                } else {
+                    0u64
+                }
+            })
+            .collect::<Vec<u64>>();
+        // reverse sort
+        epochs.sort_unstable_by(|a, b| b.cmp(a));
+
+        // move through the epochs from largest to smallest, first one that's <= ```epoch_in_question```
+        // and Bob's your uncle
+        for item in epochs {
+            if item <= epoch_in_question {
+                return Ok(item);
+            }
+        }
+        Err(StorageError::GetError(format!(
+            "Node (val: {}, len: {}) did not exist <= epoch {}",
+            node_label.val, node_label.len, epoch_in_question
+        )))
+    }
 }
 
 // ===== In-Memory database w/caching (for benchmarking) ==== //
@@ -606,4 +643,41 @@ impl Storage for AsyncInMemoryDbWithCache {
         }
         Ok(map)
     }
+
+    async fn get_epoch_lte_epoch(
+        &self,
+        node_label: crate::node_state::NodeLabel,
+        epoch_in_question: u64,
+    ) -> Result<u64, StorageError> {
+        let ids = (0..epoch_in_question)
+            .map(|epoch| crate::node_state::NodeStateKey(node_label, epoch))
+            .collect::<Vec<_>>();
+        let data = self
+            .batch_get::<crate::node_state::HistoryNodeState>(ids)
+            .await?;
+        let mut epochs = data
+            .into_iter()
+            .map(|item| {
+                if let DbRecord::HistoryNodeState(state) = item {
+                    state.key.1
+                } else {
+                    0u64
+                }
+            })
+            .collect::<Vec<u64>>();
+        // reverse sort
+        epochs.sort_unstable_by(|a, b| b.cmp(a));
+
+        // move through the epochs from largest to smallest, first one that's <= ```epoch_in_question```
+        // and Bob's your uncle
+        for item in epochs {
+            if item <= epoch_in_question {
+                return Ok(item);
+            }
+        }
+        Err(StorageError::GetError(format!(
+            "Node (val: {}, len: {}) did not exist <= epoch {}",
+            node_label.val, node_label.len, epoch_in_question
+        )))
+    }
 }
diff --git a/akd/src/storage/mod.rs b/akd/src/storage/mod.rs
@@ -82,6 +82,14 @@ pub trait Storage: Clone {
     /// Retrieve a stored record from the data layer
     async fn get<St: Storable>(&self, id: St::Key) -> Result<DbRecord, StorageError>;
 
+    /// Retrieve the last epoch <= ```epoch_in_question``` where the node with ```node_key```
+    /// was edited
+    async fn get_epoch_lte_epoch(
+        &self,
+        node_label: crate::node_state::NodeLabel,
+        epoch_in_question: u64,
+    ) -> Result<u64, StorageError>;
+
     /// Retrieve a batch of records by id
     async fn batch_get<St: Storable>(
         &self,
@@ -139,14 +147,16 @@ pub trait Storage: Clone {
     fn build_history_tree_node(
         label_val: u64,
         label_len: u32,
-        epochs: Vec<u64>,
+        birth_epoch: u64,
+        last_epoch: u64,
         parent_label_val: u64,
         parent_label_len: u32,
         node_type: u8,
     ) -> HistoryTreeNode {
         HistoryTreeNode {
             label: NodeLabel::new(label_val, label_len),
-            epochs,
+            birth_epoch,
+            last_epoch,
             parent: NodeLabel::new(parent_label_val, parent_label_len),
             node_type: NodeType::from_u8(node_type),
         }

diff --git a/akd/src/storage/tests.rs b/akd/src/storage/tests.rs
@@ -77,7 +77,8 @@ async fn test_get_and_set_item<Ns: Storage>(storage: &Ns) {
 
     let node = HistoryTreeNode {
         label: NodeLabel::new(13, 4),
-        epochs: vec![123u64, 234u64, 345u64],
+        birth_epoch: 123,
+        last_epoch: 234,
         parent: NodeLabel::new(1, 1),
         node_type: NodeType::Leaf,
     };
@@ -98,7 +99,8 @@ async fn test_get_and_set_item<Ns: Storage>(storage: &Ns) {
         assert_eq!(got_node.label, node.label);
         assert_eq!(got_node.parent, node.parent);
         assert_eq!(got_node.node_type, node.node_type);
-        assert_eq!(got_node.epochs, node.epochs);
+        assert_eq!(got_node.birth_epoch, node.birth_epoch);
+        assert_eq!(got_node.last_epoch, node.last_epoch);
     } else {
         panic!("Failed to retrieve History Tree Node");
     }

diff --git a/akd/src/tests.rs b/akd/src/tests.rs
@@ -53,7 +53,10 @@ async fn test_set_child_without_hash_at_root() -> Result<(), HistoryTreeNodeErro
         root.get_latest_epoch().unwrap_or(0) == 1,
         "Latest epochs don't match!"
     );
-    assert!(root.epochs.len() == 1, "Ask yourself some pressing questions, such as: Why are there random extra epochs in the root?");
+    assert!(
+        root.birth_epoch == root.last_epoch,
+        "How would the last epoch be different from the birth epoch without an update?"
+    );
 
     Ok(())
 }
@@ -103,7 +106,10 @@ async fn test_set_children_without_hash_at_root() -> Result<(), HistoryTreeNodeE
     }
     let latest_ep = root.get_latest_epoch();
     assert!(latest_ep.unwrap_or(0) == 1, "Latest epochs don't match!");
-    assert!(root.epochs.len() == 1, "Ask yourself some pressing questions, such as: Why are there random extra epochs in the root?");
+    assert!(
+        root.birth_epoch == root.last_epoch,
+        "How would the last epoch be different from the birth epoch without an update?"
+    );
 
     Ok(())
 }
@@ -173,7 +179,10 @@ async fn test_set_children_without_hash_multiple_at_root() -> Result<(), History
     }
     let latest_ep = root.get_latest_epoch();
     assert!(latest_ep.unwrap_or(0) == 2, "Latest epochs don't match!");
-    assert!(root.epochs.len() == 2, "Ask yourself some pressing questions, such as: Why are there random extra epochs in the root?");
+    assert!(
+        root.birth_epoch < root.last_epoch,
+        "How is the last epoch not higher than the birth epoch after an udpate?"
+    );
 
     Ok(())
 }
@@ -242,7 +251,10 @@ async fn test_get_child_at_existing_epoch_multiple_at_root() -> Result<(), Histo
     }
     let latest_ep = root.get_latest_epoch();
     assert!(latest_ep.unwrap_or(0) == 2, "Latest epochs don't match!");
-    assert!(root.epochs.len() == 2, "Ask yourself some pressing questions, such as: Why are there random extra epochs in the root?");
+    assert!(
+        root.birth_epoch < root.last_epoch,
+        "How is the last epoch not higher than the birth epoch after an udpate?"
+    );
 
     Ok(())
 }
@@ -312,7 +324,10 @@ pub async fn test_get_child_at_epoch_at_root() -> Result<(), HistoryTreeNodeErro
     }
     let latest_ep = root.get_latest_epoch();
     assert!(latest_ep.unwrap_or(0) == 4, "Latest epochs don't match!");
-    assert!(root.epochs.len() == 3, "Ask yourself some pressing questions, such as: Why are there random extra epochs in the root?");
+    assert!(
+        root.birth_epoch < root.last_epoch,
+        "How is the last epoch not higher than the birth epoch after an udpate?"
+    );
 
     Ok(())
 }