runtime: Add mkvs::OverlayTree and use it in the dispatcher

.changelog/3499.bugfix.1.md

@@ -0,0 +1 @@
+runtime: Remove incorrect Sync impl on mkvs::Tree

.changelog/3499.bugfix.2.md

@@ -0,0 +1 @@
+go/storage/mkvs: Fix edge case in overlay iterator

.changelog/3499.feature.1.md

@@ -0,0 +1 @@
+runtime: Add `mkvs::Iterator` trait and `MKVS::iter` trait method

.changelog/3499.feature.2.md

@@ -0,0 +1 @@
+runtime: Add mkvs::OverlayTree similar to its Go counterpart

.changelog/3499.feature.3.md

@@ -0,0 +1 @@
+go/storage/mkvs: Generate proofs even if Tree.Position is not found

client/src/transaction/snapshot.rs

@@ -1,5 +1,5 @@
 //! A block snapshot.
-use std::any::Any;
+use std::{any::Any, sync::Mutex};
 
 use anyhow::{Context as AnyContext, Result};
 use grpcio::CallOption;
@@ -11,7 +11,7 @@ use oasis_core_runtime::{
         roothash::{Block, Namespace},
     },
     storage::{
-        mkvs::{sync::*, Prefix, Root, Tree, WriteLog},
+        mkvs::{sync::*, Iterator, Prefix, Root, Tree, WriteLog},
         MKVS,
     },
     transaction::types::{TxnCall, TxnOutput},
@@ -57,7 +57,7 @@ pub struct BlockSnapshot {
     pub block_hash: Hash,
 
     read_syncer: RemoteReadSync,
-    mkvs: Tree,
+    mkvs: Mutex<Tree>,
 }
 
 impl Clone for BlockSnapshot {
@@ -77,7 +77,7 @@ impl Clone for BlockSnapshot {
             block,
             block_hash,
             read_syncer,
-            mkvs,
+            mkvs: Mutex::new(mkvs),
         }
     }
 }
@@ -97,18 +97,20 @@ impl BlockSnapshot {
             block_hash: block.header.encoded_hash(),
             block,
             read_syncer,
-            mkvs,
+            mkvs: Mutex::new(mkvs),
         }
     }
 }
 
 impl MKVS for BlockSnapshot {
     fn get(&self, ctx: Context, key: &[u8]) -> Option<Vec<u8>> {
-        MKVS::get(&self.mkvs, ctx, key)
+        let mkvs = self.mkvs.lock().unwrap();
+        mkvs.get(ctx, key).unwrap()
     }
 
     fn cache_contains_key(&self, ctx: Context, key: &[u8]) -> bool {
-        MKVS::cache_contains_key(&self.mkvs, ctx, key)
+        let mkvs = self.mkvs.lock().unwrap();
+        mkvs.cache_contains_key(ctx, key)
     }
 
     fn insert(&mut self, _ctx: Context, _key: &[u8], _value: &[u8]) -> Option<Vec<u8>> {
@@ -120,7 +122,12 @@ impl MKVS for BlockSnapshot {
     }
 
     fn prefetch_prefixes(&self, ctx: Context, prefixes: &Vec<Prefix>, limit: u16) {
-        MKVS::prefetch_prefixes(&self.mkvs, ctx, prefixes, limit)
+        let mkvs = self.mkvs.lock().unwrap();
+        mkvs.prefetch_prefixes(ctx, prefixes, limit).unwrap()
+    }
+
+    fn iter(&self, _ctx: Context) -> Box<dyn Iterator + '_> {
+        unimplemented!("block snapshot doesn't support iterators");
     }
 
     fn commit(
@@ -131,10 +138,6 @@ impl MKVS for BlockSnapshot {
     ) -> Result<(WriteLog, Hash)> {
         unimplemented!("block snapshot is read-only");
     }
-
-    fn rollback(&mut self) {
-        unimplemented!("block snapshot is read-only");
-    }
 }
 
 #[derive(Clone)]

go/storage/mkvs/iterator.go

@@ -150,7 +150,7 @@ func IteratorPrefetch(prefetch uint16) IteratorOption {
 // visited nodes.
 func WithProof(root hash.Hash) IteratorOption {
 	return func(it Iterator) {
-		it.(*treeIterator).proofBuilder = syncer.NewProofBuilder(root)
+		it.(*treeIterator).proofBuilder = syncer.NewProofBuilder(root, root)
 	}
 }
 
@@ -251,10 +251,7 @@ func (it *treeIterator) doNext(ptr *node.Pointer, bitDepth node.Depth, path, key
 
 	// Include nodes in proof if we have a proof builder.
 	if pb := it.proofBuilder; pb != nil && ptr != nil {
-		proofRoot := pb.GetRoot()
-		if pb.HasRoot() || proofRoot.Equal(&ptr.Hash) {
-			pb.Include(nd)
-		}
+		pb.Include(nd)
 	}
 
 	switch n := nd.(type) {

go/storage/mkvs/lookup.go

@@ -48,7 +48,7 @@ func (t *tree) SyncGet(ctx context.Context, request *syncer.GetRequest) (*syncer
 	// Remember where the path from root to target node ends (will end).
 	t.cache.markPosition()
 
-	pb := syncer.NewProofBuilder(request.Tree.Position)
+	pb := syncer.NewProofBuilder(request.Tree.Root.Hash, request.Tree.Position)
 	opts := doGetOptions{
 		proofBuilder:    pb,
 		includeSiblings: request.IncludeSiblings,
@@ -108,10 +108,7 @@ func (t *tree) doGet(
 
 	// Include nodes in proof if we have a proof builder.
 	if pb := opts.proofBuilder; pb != nil && ptr != nil {
-		proofRoot := pb.GetRoot()
-		if pb.HasRoot() || proofRoot.Equal(&ptr.Hash) {
-			pb.Include(nd)
-		}
+		pb.Include(nd)
 	}
 
 	// This may be used to only include the given node in a proof and not

go/storage/mkvs/overlay.go

@@ -111,6 +111,7 @@ func (o *treeOverlay) Commit(ctx context.Context) error {
 			return err
 		}
 	}
+	o.dirty = make(map[string]bool)
 
 	return nil
 }
@@ -169,7 +170,7 @@ func (it *treeOverlayIterator) Seek(key node.Key) {
 }
 
 func (it *treeOverlayIterator) Next() {
-	if !it.overlay.Valid() || it.inner.Key().Compare(it.overlay.Key()) <= 0 {
+	if !it.overlay.Valid() || (it.inner.Valid() && it.inner.Key().Compare(it.overlay.Key()) <= 0) {
 		// Key of inner iterator is smaller or equal than the key of the overlay iterator.
 		it.inner.Next()
 	} else {

go/storage/mkvs/overlay_test.go

@@ -12,12 +12,9 @@ import (
 
 func TestOverlay(t *testing.T) {
 	require := require.New(t)
-
 	ctx := context.Background()
-	tree := New(nil, nil)
-	defer tree.Close()
 
-	// Insert some items.
+	// Generate some items.
 	items := writelog.WriteLog{
 		writelog.LogEntry{Key: []byte("key"), Value: []byte("first")},
 		writelog.LogEntry{Key: []byte("key 1"), Value: []byte("one")},
@@ -40,11 +37,30 @@ func TestOverlay(t *testing.T) {
 		{seek: node.Key("key A"), pos: -1},
 	}
 
+	tree := New(nil, nil)
+	defer tree.Close()
+
+	// Create an overlay over an empty tree and insert some items into the overlay.
+	overlay := NewOverlay(tree)
+	for _, item := range items {
+		err := overlay.Insert(ctx, item.Key, item.Value)
+		require.NoError(err, "Insert")
+	}
+
+	// Test that an overlay-only iterator works correctly.
+	t.Run("OnlyOverlay/Iterator", func(t *testing.T) {
+		it := overlay.NewIterator(ctx)
+		defer it.Close()
+
+		testIterator(t, items, it, tests)
+	})
+
+	// Insert some items into the underlying tree.
 	err := tree.ApplyWriteLog(ctx, writelog.NewStaticIterator(items))
 	require.NoError(err, "ApplyWriteLog")
 
 	// Create an overlay.
-	overlay := NewOverlay(tree)
+	overlay = NewOverlay(tree)
 
 	// Test that all keys can be fetched from an empty overlay.
 	t.Run("EmptyOverlay/Get", func(t *testing.T) {

go/storage/mkvs/syncer/proof.go

@@ -34,14 +34,16 @@ type proofNode struct {
 // ProofBuilder is a Merkle proof builder.
 type ProofBuilder struct {
 	root     hash.Hash
+	subtree  hash.Hash
 	included map[hash.Hash]*proofNode
 	size     uint64
 }
 
 // NewProofBuilder creates a new Merkle proof builder for the given root.
-func NewProofBuilder(root hash.Hash) *ProofBuilder {
+func NewProofBuilder(root, subtree hash.Hash) *ProofBuilder {
 	return &ProofBuilder{
 		root:     root,
+		subtree:  subtree,
 		included: make(map[hash.Hash]*proofNode),
 	}
 }
@@ -94,14 +96,14 @@ func (b *ProofBuilder) Include(n node.Node) {
 	b.size += 1 + uint64(len(pn.serialized))
 }
 
-// HasRoot returns true if the root node has already been included.
-func (b *ProofBuilder) HasRoot() bool {
-	return b.included[b.root] != nil
+// HasSubtree returns true if the subtree root node has already been included.
+func (b *ProofBuilder) HasSubtreeRoot() bool {
+	return b.included[b.subtree] != nil
 }
 
-// GetRoot returns the root hash for this proof.
-func (b *ProofBuilder) GetRoot() hash.Hash {
-	return b.root
+// GetSubtree returns the subtree root hash for this proof.
+func (b *ProofBuilder) GetSubtreeRoot() hash.Hash {
+	return b.subtree
 }
 
 // Size returns the current size of this proof.
@@ -111,10 +113,17 @@ func (b *ProofBuilder) Size() uint64 {
 
 // Build tries to build the proof.
 func (b *ProofBuilder) Build(ctx context.Context) (*Proof, error) {
-	proof := Proof{
-		UntrustedRoot: b.root,
+	var proof Proof
+	switch b.HasSubtreeRoot() {
+	case true:
+		// A partial proof for the subtree is available, include that.
+		proof.UntrustedRoot = b.subtree
+	case false:
+		// No partial proof available, we need to use the tree root.
+		proof.UntrustedRoot = b.root
 	}
-	if err := b.build(ctx, &proof, b.root); err != nil {
+
+	if err := b.build(ctx, &proof, proof.UntrustedRoot); err != nil {
 		return nil, err
 	}
 	return &proof, nil

go/storage/mkvs/syncer_test.go

@@ -31,9 +31,9 @@ func TestProof(t *testing.T) {
 	require.NoError(err, "Commit")
 
 	// Create a Merkle proof, starting at the root node.
-	builder := syncer.NewProofBuilder(rootHash)
-	require.False(builder.HasRoot(), "HasRoot should return false")
-	require.EqualValues(rootHash, builder.GetRoot(), "GetRoot should return correct root")
+	builder := syncer.NewProofBuilder(rootHash, rootHash)
+	require.False(builder.HasSubtreeRoot(), "HasSubtreeRoot should return false")
+	require.EqualValues(rootHash, builder.GetSubtreeRoot(), "GetSubtreeRoot should return correct root")
 
 	rootOnlyProof, err := builder.Build(ctx)
 	require.NoError(err, "Build should not fail without a root present")
@@ -44,7 +44,7 @@ func TestProof(t *testing.T) {
 	// Include root node.
 	rootNode := tree.cache.pendingRoot.Node
 	builder.Include(rootNode)
-	require.True(builder.HasRoot(), "HasRoot should return true after root included")
+	require.True(builder.HasSubtreeRoot(), "HasRoot should return true after root included")
 
 	proof, err := builder.Build(ctx)
 	require.NoError(err, "Build should not fail")
@@ -107,7 +107,7 @@ func TestProof(t *testing.T) {
 	// Empty root proof should verify.
 	var emptyHash hash.Hash
 	emptyHash.Empty()
-	builder = syncer.NewProofBuilder(emptyHash)
+	builder = syncer.NewProofBuilder(emptyHash, emptyHash)
 	emptyRootProof, err := builder.Build(ctx)
 	require.NoError(err, "Build should not fail for an empty root")
 	emptyRootPtr, err := pv.VerifyProof(ctx, emptyHash, emptyRootProof)

go/storage/mkvs/tree_test.go

@@ -1634,6 +1634,46 @@ func testPruneLoneRootsShared2(t *testing.T, ndb db.NodeDB, factory NodeDBFactor
 	require.NoError(t, it.Err(), "tree should still be consistent")
 }
 
+func testPruneLoneRootsShared3(t *testing.T, ndb db.NodeDB, factory NodeDBFactory) {
+	require := require.New(t)
+	ctx := context.Background()
+
+	// Create a root in version 0.
+	tree := New(nil, ndb)
+	err := tree.Insert(ctx, []byte("foo"), []byte("bar"))
+	require.NoError(err, "Insert")
+	_, _, err = tree.Commit(ctx, testNs, 0)
+	require.NoError(err, "Commit")
+
+	// Create another root in version 0.
+	tree = New(nil, ndb)
+	err = tree.Insert(ctx, []byte("moo"), []byte("goo"))
+	require.NoError(err, "Insert")
+	_, rootHashR0_2, err := tree.Commit(ctx, testNs, 0)
+	require.NoError(err, "Commit")
+
+	// Create the same root as the first root in version 1.
+	tree = New(nil, ndb)
+	err = tree.Insert(ctx, []byte("foo"), []byte("bar"))
+	require.NoError(err, "Insert")
+	_, rootHashR1_1, err := tree.Commit(ctx, testNs, 1)
+	require.NoError(err, "Commit")
+
+	// Finalize version 0 with the second root.
+	err = ndb.Finalize(ctx, 0, []hash.Hash{rootHashR0_2})
+	require.NoError(err, "Finalize")
+
+	// Make sure that the first root in version 1 is still valid.
+	tree = NewWithRoot(nil, ndb, node.Root{
+		Namespace: testNs,
+		Version:   1,
+		Hash:      rootHashR1_1,
+	})
+	value, err := tree.Get(ctx, []byte("foo"))
+	require.NoError(err, "Get")
+	require.EqualValues([]byte("bar"), value)
+}
+
 func testPruneLoneRoots(t *testing.T, ndb db.NodeDB, factory NodeDBFactory) {
 	ctx := context.Background()
 
@@ -2064,6 +2104,7 @@ func testBackend(
 		{"PruneLoneRoots", testPruneLoneRoots},
 		{"PruneLoneRootsShared", testPruneLoneRootsShared},
 		{"PruneLoneRootsShared2", testPruneLoneRootsShared2},
+		{"PruneLoneRootsShared3", testPruneLoneRootsShared3},
 		{"PruneForkedRoots", testPruneForkedRoots},
 		{"SpecialCase1", testSpecialCase1},
 		{"SpecialCase2", testSpecialCase2},