merkledb -- dynamic root

scripts/mocks.mockgen.txt

@@ -45,5 +45,4 @@ github.com/ava-labs/avalanchego/vms/registry=VMGetter=vms/registry/mock_vm_gette
 github.com/ava-labs/avalanchego/vms/registry=VMRegisterer=vms/registry/mock_vm_registerer.go
 github.com/ava-labs/avalanchego/vms/registry=VMRegistry=vms/registry/mock_vm_registry.go
 github.com/ava-labs/avalanchego/vms=Factory,Manager=vms/mock_manager.go
-github.com/ava-labs/avalanchego/x/merkledb=MerkleDB=x/merkledb/mock_db.go
 github.com/ava-labs/avalanchego/x/sync=Client=x/sync/mock_client.go

x/merkledb/codec.go

@@ -66,11 +66,13 @@ type encoder interface {
 	// Returns the bytes that will be hashed to generate [n]'s ID.
 	// Assumes [n] is non-nil.
 	encodeHashValues(n *node) []byte
+	encodeKey(key Key) []byte
 }
 
 type decoder interface {
 	// Assumes [n] is non-nil.
 	decodeDBNode(bytes []byte, n *dbNode) error
+	decodeKey(bytes []byte) (Key, error)
 }
 
 func newCodec() encoderDecoder {
@@ -98,7 +100,6 @@ func (c *codecImpl) encodeDBNode(n *dbNode) []byte {
 		estimatedLen = estimatedValueLen + minVarIntLen + estimatedNodeChildLen*numChildren
 		buf          = bytes.NewBuffer(make([]byte, 0, estimatedLen))
 	)
-
 	c.encodeMaybeByteSlice(buf, n.value)
 	c.encodeUint(buf, uint64(numChildren))
 	// Note we insert children in order of increasing index
@@ -108,7 +109,7 @@ func (c *codecImpl) encodeDBNode(n *dbNode) []byte {
 	for _, index := range keys {
 		entry := n.children[index]
 		c.encodeUint(buf, uint64(index))
-		c.encodeKey(buf, entry.compressedKey)
+		c.encodeKeyToBuffer(buf, entry.compressedKey)
 		_, _ = buf.Write(entry.id[:])
 		c.encodeBool(buf, entry.hasValue)
 	}
@@ -134,7 +135,7 @@ func (c *codecImpl) encodeHashValues(n *node) []byte {
 		_, _ = buf.Write(entry.id[:])
 	}
 	c.encodeMaybeByteSlice(buf, n.valueDigest)
-	c.encodeKey(buf, n.key)
+	c.encodeKeyToBuffer(buf, n.key)
 
 	return buf.Bytes()
 }
@@ -172,7 +173,7 @@ func (c *codecImpl) decodeDBNode(b []byte, n *dbNode) error {
 		}
 		previousChild = index
 
-		compressedKey, err := c.decodeKey(src)
+		compressedKey, err := c.decodeKeyFromReader(src)
 		if err != nil {
 			return err
 		}
@@ -330,12 +331,31 @@ func (*codecImpl) decodeID(src *bytes.Reader) (ids.ID, error) {
 	return id, err
 }
 
-func (c *codecImpl) encodeKey(dst *bytes.Buffer, key Key) {
+func (c *codecImpl) encodeKey(key Key) []byte {
+	estimatedLen := binary.MaxVarintLen64 + len(key.Bytes())
+	dst := bytes.NewBuffer(make([]byte, 0, estimatedLen))
+	c.encodeKeyToBuffer(dst, key)
+	return dst.Bytes()
+}
+
+func (c *codecImpl) encodeKeyToBuffer(dst *bytes.Buffer, key Key) {
 	c.encodeUint(dst, uint64(key.length))
 	_, _ = dst.Write(key.Bytes())
 }
 
-func (c *codecImpl) decodeKey(src *bytes.Reader) (Key, error) {
+func (c *codecImpl) decodeKey(b []byte) (Key, error) {
+	src := bytes.NewReader(b)
+	key, err := c.decodeKeyFromReader(src)
+	if err != nil {
+		return Key{}, err
+	}
+	if src.Len() != 0 {
+		return Key{}, errExtraSpace
+	}
+	return key, err
+}
+
+func (c *codecImpl) decodeKeyFromReader(src *bytes.Reader) (Key, error) {
 	if minKeyLen > src.Len() {
 		return Key{}, io.ErrUnexpectedEOF
 	}

x/merkledb/codec_test.go

@@ -81,21 +81,14 @@ func FuzzCodecKey(f *testing.F) {
 		) {
 			require := require.New(t)
 			codec := codec.(*codecImpl)
-			reader := bytes.NewReader(b)
-			startLen := reader.Len()
-			got, err := codec.decodeKey(reader)
+			got, err := codec.decodeKey(b)
 			if err != nil {
 				t.SkipNow()
 			}
-			endLen := reader.Len()
-			numRead := startLen - endLen
 
 			// Encoding [got] should be the same as [b].
-			var buf bytes.Buffer
-			codec.encodeKey(&buf, got)
-			bufBytes := buf.Bytes()
-			require.Len(bufBytes, numRead)
-			require.Equal(b[:numRead], bufBytes)
+			gotBytes := codec.encodeKey(got)
+			require.Equal(b, gotBytes)
 		},
 	)
 }
@@ -248,7 +241,6 @@ func FuzzEncodeHashValues(f *testing.F) {
 
 func TestCodecDecodeKeyLengthOverflowRegression(t *testing.T) {
 	codec := codec.(*codecImpl)
-	bytes := bytes.NewReader(binary.AppendUvarint(nil, math.MaxInt))
-	_, err := codec.decodeKey(bytes)
+	_, err := codec.decodeKey(binary.AppendUvarint(nil, math.MaxInt))
 	require.ErrorIs(t, err, io.ErrUnexpectedEOF)
 }

x/merkledb/db.go

@@ -51,11 +51,11 @@ var (
 	intermediateNodePrefix = []byte{2}
 
 	cleanShutdownKey        = []byte(string(metadataPrefix) + "cleanShutdown")
+	rootDBKey               = []byte(string(metadataPrefix) + "root")
 	hadCleanShutdown        = []byte{1}
 	didNotHaveCleanShutdown = []byte{0}
 
-	errSameRoot      = errors.New("start and end root are the same")
-	errNoNewSentinel = errors.New("there was no updated sentinel node in change list")
+	errSameRoot = errors.New("start and end root are the same")
 )
 
 type ChangeProofer interface {
@@ -64,6 +64,9 @@ type ChangeProofer interface {
 	// Returns at most [maxLength] key/value pairs.
 	// Returns [ErrInsufficientHistory] if this node has insufficient history
 	// to generate the proof.
+	// Returns ErrEmptyProof if [endRootID] is ids.Empty.
+	// Note that [endRootID] == ids.Empty means the trie is empty
+	// (i.e. we don't need a change proof.)
 	GetChangeProof(
 		ctx context.Context,
 		startRootID ids.ID,
@@ -100,6 +103,9 @@ type RangeProofer interface {
 	// [start, end] when the root of the trie was [rootID].
 	// If [start] is Nothing, there's no lower bound on the range.
 	// If [end] is Nothing, there's no upper bound on the range.
+	// Returns ErrEmptyProof if [rootID] is ids.Empty.
+	// Note that [rootID] == ids.Empty means the trie is empty
+	// (i.e. we don't need a range proof.)
 	GetRangeProofAtRoot(
 		ctx context.Context,
 		rootID ids.ID,
@@ -201,11 +207,11 @@ type merkleDB struct {
 	debugTracer trace.Tracer
 	infoTracer  trace.Tracer
 
-	// The sentinel node of this trie.
-	// It is the node with a nil key and is the ancestor of all nodes in the trie.
-	// If it has a value or has multiple children, it is also the root of the trie.
-	sentinelNode *node
-	rootID       ids.ID
+	// The root of this trie.
+	// Nothing if the trie is empty.
+	root maybe.Maybe[*node]
+
+	rootID ids.ID
 
 	// Valid children of this trie.
 	childViews []*trieView
@@ -268,6 +274,9 @@ func newDatabase(
 	// add current root to history (has no changes)
 	trieDB.history.record(&changeSummary{
 		rootID: trieDB.rootID,
+		rootChange: change[maybe.Maybe[*node]]{
+			after: trieDB.root,
+		},
 		values: map[Key]*change[maybe.Maybe[[]byte]]{},
 		nodes:  map[Key]*change[*node]{},
 	})
@@ -295,7 +304,8 @@ func newDatabase(
 // Deletes every intermediate node and rebuilds them by re-adding every key/value.
 // TODO: make this more efficient by only clearing out the stale portions of the trie.
 func (db *merkleDB) rebuild(ctx context.Context, cacheSize int) error {
-	db.sentinelNode = newNode(Key{})
+	db.root = maybe.Nothing[*node]()
+	db.rootID = ids.Empty
 
 	// Delete intermediate nodes.
 	if err := database.ClearPrefix(db.baseDB, intermediateNodePrefix, rebuildIntermediateDeletionWriteSize); err != nil {
@@ -581,13 +591,6 @@ func (db *merkleDB) getMerkleRoot() ids.ID {
 	return db.rootID
 }
 
-// isSentinelNodeTheRoot returns true if the passed in sentinel node has a value and or multiple child nodes
-// When this is true, the root of the trie is the sentinel node
-// When this is false, the root of the trie is the sentinel node's single child
-func isSentinelNodeTheRoot(sentinel *node) bool {
-	return sentinel.valueDigest.HasValue() || len(sentinel.children) != 1
-}
-
 func (db *merkleDB) GetProof(ctx context.Context, key []byte) (*Proof, error) {
 	db.commitLock.RLock()
 	defer db.commitLock.RUnlock()
@@ -644,11 +647,13 @@ func (db *merkleDB) getRangeProofAtRoot(
 	end maybe.Maybe[[]byte],
 	maxLength int,
 ) (*RangeProof, error) {
-	if db.closed {
+	switch {
+	case db.closed:
 		return nil, database.ErrClosed
-	}
-	if maxLength <= 0 {
+	case maxLength <= 0:
 		return nil, fmt.Errorf("%w but was %d", ErrInvalidMaxLength, maxLength)
+	case rootID == ids.Empty:
+		return nil, ErrEmptyProof
 	}
 
 	historicalView, err := db.getHistoricalViewForRange(rootID, start, end)
@@ -666,11 +671,13 @@ func (db *merkleDB) GetChangeProof(
 	end maybe.Maybe[[]byte],
 	maxLength int,
 ) (*ChangeProof, error) {
-	if start.HasValue() && end.HasValue() && bytes.Compare(start.Value(), end.Value()) == 1 {
+	switch {
+	case start.HasValue() && end.HasValue() && bytes.Compare(start.Value(), end.Value()) == 1:
 		return nil, ErrStartAfterEnd
-	}
-	if startRootID == endRootID {
+	case startRootID == endRootID:
 		return nil, errSameRoot
+	case endRootID == ids.Empty:
+		return nil, ErrEmptyProof
 	}
 
 	db.commitLock.RLock()
@@ -931,13 +938,7 @@ func (db *merkleDB) commitChanges(ctx context.Context, trieToCommit *trieView) e
 		return nil
 	}
 
-	sentinelChange, ok := changes.nodes[Key{}]
-	if !ok {
-		return errNoNewSentinel
-	}
-
 	currentValueNodeBatch := db.valueNodeDB.NewBatch()
-
 	_, nodesSpan := db.infoTracer.Start(ctx, "MerkleDB.commitChanges.writeNodes")
 	for key, nodeChange := range changes.nodes {
 		shouldAddIntermediate := nodeChange.after != nil && !nodeChange.after.hasValue()
@@ -973,12 +974,18 @@ func (db *merkleDB) commitChanges(ctx context.Context, trieToCommit *trieView) e
 		return err
 	}
 
-	// Only modify in-memory state after the commit succeeds
-	// so that we don't need to clean up on error.
-	db.sentinelNode = sentinelChange.after
-	db.rootID = changes.rootID
 	db.history.record(changes)
-	return nil
+
+	// Update root in database.
+	db.root = changes.rootChange.after
+	db.rootID = changes.rootID
+
+	if db.root.IsNothing() {
+		return db.baseDB.Delete(rootDBKey)
+	}
+
+	rootKey := codec.encodeKey(db.root.Value().key)
+	return db.baseDB.Put(rootDBKey, rootKey)
 }
 
 // moveChildViewsToDB removes any child views from the trieToCommit and moves them to the db
@@ -1014,7 +1021,7 @@ func (db *merkleDB) VerifyChangeProof(
 	case start.HasValue() && end.HasValue() && bytes.Compare(start.Value(), end.Value()) > 0:
 		return ErrStartAfterEnd
 	case proof.Empty():
-		return ErrNoMerkleProof
+		return ErrEmptyProof
 	case end.HasValue() && len(proof.KeyChanges) == 0 && len(proof.EndProof) == 0:
 		// We requested an end proof but didn't get one.
 		return ErrNoEndProof
@@ -1156,37 +1163,41 @@ func (db *merkleDB) invalidateChildrenExcept(exception *trieView) {
 	}
 }
 
+// If the root is on disk, set [db.root] to it.
+// Otherwise leave [db.root] as Nothing.
 func (db *merkleDB) initializeRoot() error {
-	// Not sure if the  sentinel node exists or if it had a value,
-	// so check under both prefixes
-	var err error
-	db.sentinelNode, err = db.intermediateNodeDB.Get(Key{})
+	rootKeyBytes, err := db.baseDB.Get(rootDBKey)
+	if err != nil {
+		if !errors.Is(err, database.ErrNotFound) {
+			return err
+		}
+		// Root isn't on disk.
+		return nil
+	}
 
-	if errors.Is(err, database.ErrNotFound) {
-		// Didn't find the sentinel in the intermediateNodeDB, check the valueNodeDB
-		db.sentinelNode, err = db.valueNodeDB.Get(Key{})
+	// Root is on disk.
+	rootKey, err := codec.decodeKey(rootKeyBytes)
+	if err != nil {
+		return err
 	}
 
+	// First, see if root is an intermediate node.
+	var root *node
+	root, err = db.getEditableNode(rootKey, false /* hasValue */)
 	if err != nil {
 		if !errors.Is(err, database.ErrNotFound) {
 			return err
 		}
 
-		// Sentinel node doesn't exist in either database prefix.
-		// Make a new one and store it in the intermediateNodeDB
-		db.sentinelNode = newNode(Key{})
-		if err := db.intermediateNodeDB.Put(Key{}, db.sentinelNode); err != nil {
+		// The root must be a value node.
+		root, err = db.getEditableNode(rootKey, true /* hasValue */)
+		if err != nil {
 			return err
 		}
 	}
 
-	db.rootID = db.sentinelNode.calculateID(db.metrics)
-	if !isSentinelNodeTheRoot(db.sentinelNode) {
-		// If the sentinel node is not the root, the trie's root is the sentinel node's only child
-		for _, childEntry := range db.sentinelNode.children {
-			db.rootID = childEntry.id
-		}
-	}
+	db.rootID = root.calculateID(db.metrics)
+	db.root = maybe.Some(root)
 	return nil
 }
 
@@ -1263,12 +1274,15 @@ func (db *merkleDB) getNode(key Key, hasValue bool) (*node, error) {
 	switch {
 	case db.closed:
 		return nil, database.ErrClosed
-	case key == Key{}:
-		return db.sentinelNode, nil
 	case hasValue:
 		return db.valueNodeDB.Get(key)
+	default:
+		return db.intermediateNodeDB.Get(key)
 	}
-	return db.intermediateNodeDB.Get(key)
+}
+
+func (db *merkleDB) getRoot() maybe.Maybe[*node] {
+	return db.root
 }
 
 func (db *merkleDB) Clear() error {
@@ -1287,13 +1301,13 @@ func (db *merkleDB) Clear() error {
 	}
 
 	// Clear root
-	db.sentinelNode = newNode(Key{})
-	db.rootID = db.sentinelNode.calculateID(db.metrics)
+	db.root = maybe.Nothing[*node]()
+	db.rootID = ids.Empty
 
 	// Clear history
 	db.history = newTrieHistory(db.history.maxHistoryLen)
 	db.history.record(&changeSummary{
-		rootID: db.getMerkleRoot(),
+		rootID: db.rootID,
 		values: map[Key]*change[maybe.Maybe[[]byte]]{},
 		nodes:  map[Key]*change[*node]{},
 	})