runtime/storage/mkvs: Add method for checking local key existence

.changelog/2938.feature.md

@@ -0,0 +1,4 @@
+runtime/storage/mkvs: Add method for checking local key existence
+
+Adds a method to probe the local cache for key existence, guaranteeing
+that no remote syncing will be done.

client/src/transaction/snapshot.rs

@@ -107,6 +107,10 @@ impl MKVS for BlockSnapshot {
         MKVS::get(&self.mkvs, ctx, key)
     }
 
+    fn cache_contains_key(&self, ctx: Context, key: &[u8]) -> bool {
+        MKVS::cache_contains_key(&self.mkvs, ctx, key)
+    }
+
     fn insert(&mut self, _ctx: Context, _key: &[u8], _value: &[u8]) -> Option<Vec<u8>> {
         unimplemented!("block snapshot is read-only");
     }

runtime/src/storage/mkvs/cache/cache.rs

@@ -79,11 +79,13 @@ pub trait Cache {
     /// Dereference a node pointer into a concrete node object.
     ///
     /// Calling this method may invoke the underlying read syncer.
+    /// Giving a None fetcher forces the dereference to be local-only,
+    /// without invoking the read syncer.
     fn deref_node_ptr<F: ReadSyncFetcher>(
         &mut self,
         ctx: &Arc<Context>,
         ptr: NodePtrRef,
-        fetcher: F,
+        fetcher: Option<F>,
     ) -> Fallible<Option<NodeRef>>;
     /// Perform a remote sync with the configured remote syncer.
     fn remote_sync<F: ReadSyncFetcher>(

runtime/src/storage/mkvs/cache/lru_cache.rs

@@ -415,7 +415,7 @@ impl Cache for LRUCache {
         &mut self,
         ctx: &Arc<Context>,
         ptr: NodePtrRef,
-        fetcher: F,
+        fetcher: Option<F>,
     ) -> Fallible<Option<NodeRef>> {
         let ptr_ref = ptr;
         let ptr = ptr_ref.borrow();
@@ -447,7 +447,13 @@ impl Cache for LRUCache {
         }
 
         // Node not available locally, fetch from read syncer.
-        self.remote_sync(ctx, ptr_ref.clone(), fetcher)?;
+        if let Some(fetcher) = fetcher {
+            self.remote_sync(ctx, ptr_ref.clone(), fetcher)?;
+        } else {
+            return Err(format_err!(
+                "mkvs: node to dereference not available locally and no fetcher provided"
+            ));
+        }
 
         let ptr = ptr_ref.borrow();
         if ptr.node.is_none() {

runtime/src/storage/mkvs/mod.rs

@@ -120,6 +120,13 @@ pub trait MKVS: Send + Sync {
     /// Fetch entry with given key.
     fn get(&self, ctx: Context, key: &[u8]) -> Option<Vec<u8>>;
 
+    /// Check if the local MKVS cache contains the given key.
+    ///
+    /// While get can be used to check if the MKVS as a whole contains
+    /// a given key, this function specifically guarantees that no remote
+    /// syncing will be invoked, only checking the local cache.
+    fn cache_contains_key(&self, ctx: Context, key: &[u8]) -> bool;
+
     /// Update entry with given key.
     ///
     /// If the database did not have this key present, [`None`] is returned.

runtime/src/storage/mkvs/tree/insert.rs

@@ -53,7 +53,7 @@ impl Tree {
         let node_ref = self.cache.borrow_mut().deref_node_ptr(
             ctx,
             ptr.clone(),
-            FetcherSyncGet::new(key, false),
+            Some(FetcherSyncGet::new(key, false)),
         )?;
 
         let (_, key_remainder) = key.split(bit_depth, key.bit_length());

runtime/src/storage/mkvs/tree/iterator.rs

@@ -195,7 +195,7 @@ impl<'tree> TreeIterator<'tree> {
         let node_ref = self.tree.cache.borrow_mut().deref_node_ptr(
             &self.ctx,
             ptr.clone(),
-            FetcherSyncIterate::new(&key, self.prefetch),
+            Some(FetcherSyncIterate::new(&key, self.prefetch)),
         )?;
 
         match classify_noderef!(?node_ref) {

runtime/src/storage/mkvs/tree/lookup.rs

@@ -45,6 +45,19 @@ impl<'a> ReadSyncFetcher for FetcherSyncGet<'a> {
 impl Tree {
     /// Get an existing key.
     pub fn get(&self, ctx: Context, key: &[u8]) -> Fallible<Option<Vec<u8>>> {
+        self._get_top(ctx, key, false)
+    }
+
+    /// Check if the key exists in the local cache.
+    pub fn cache_contains_key(&self, ctx: Context, key: &[u8]) -> bool {
+        match self._get_top(ctx, key, true) {
+            Ok(Some(_)) => true,
+            Ok(None) => false,
+            Err(_) => false,
+        }
+    }
+
+    fn _get_top(&self, ctx: Context, key: &[u8], check_only: bool) -> Fallible<Option<Vec<u8>>> {
         let ctx = ctx.freeze();
         let boxed_key = key.to_vec();
         let pending_root = self.cache.borrow().get_pending_root();
@@ -57,7 +70,7 @@ impl Tree {
         // Remember where the path from root to target node ends (will end).
         self.cache.borrow_mut().mark_position();
 
-        Ok(self._get(&ctx, pending_root, 0, &boxed_key, 0)?)
+        Ok(self._get(&ctx, pending_root, 0, &boxed_key, 0, check_only)?)
     }
 
     fn _get(
@@ -67,11 +80,17 @@ impl Tree {
         bit_depth: Depth,
         key: &Key,
         depth: Depth,
+        check_only: bool,
     ) -> Fallible<Option<Value>> {
-        let node_ref =
-            self.cache
-                .borrow_mut()
-                .deref_node_ptr(ctx, ptr, FetcherSyncGet::new(key, false))?;
+        let node_ref = self.cache.borrow_mut().deref_node_ptr(
+            ctx,
+            ptr,
+            if check_only {
+                None
+            } else {
+                Some(FetcherSyncGet::new(key, false))
+            },
+        )?;
 
         match classify_noderef!(?node_ref) {
             NodeKind::None => {
@@ -90,6 +109,7 @@ impl Tree {
                             bit_depth + n.label_bit_length,
                             key,
                             depth,
+                            check_only,
                         );
                     }
 
@@ -106,6 +126,7 @@ impl Tree {
                             bit_depth + n.label_bit_length,
                             key,
                             depth + 1,
+                            check_only,
                         );
                     } else {
                         return self._get(
@@ -114,6 +135,7 @@ impl Tree {
                             bit_depth + n.label_bit_length,
                             key,
                             depth + 1,
+                            check_only,
                         );
                     }
                 }

runtime/src/storage/mkvs/tree/mkvs.rs

@@ -17,6 +17,11 @@ impl MKVS for Tree {
         self.get(ctx, key).unwrap()
     }
 
+    fn cache_contains_key(&self, ctx: Context, key: &[u8]) -> bool {
+        let _lock = self.lock.lock().unwrap();
+        self.cache_contains_key(ctx, key)
+    }
+
     fn insert(&mut self, ctx: Context, key: &[u8], value: &[u8]) -> Option<Vec<u8>> {
         let lock = self.lock.clone();
         let _guard = lock.lock().unwrap();

runtime/src/storage/mkvs/tree/remove.rs

@@ -54,7 +54,7 @@ impl Tree {
         let node_ref = self.cache.borrow_mut().deref_node_ptr(
             ctx,
             ptr.clone(),
-            FetcherSyncGet::new(key, true),
+            Some(FetcherSyncGet::new(key, true)),
         )?;
 
         match classify_noderef!(?node_ref) {
@@ -106,12 +106,12 @@ impl Tree {
                     remaining_left = self.cache.borrow_mut().deref_node_ptr(
                         ctx,
                         n.left.clone(),
-                        FetcherSyncGet::new(key, true),
+                        Some(FetcherSyncGet::new(key, true)),
                     )?;
                     remaining_right = self.cache.borrow_mut().deref_node_ptr(
                         ctx,
                         n.right.clone(),
-                        FetcherSyncGet::new(key, true),
+                        Some(FetcherSyncGet::new(key, true)),
                     )?;
                 } else {
                     unreachable!("node kind is Internal");