Skip to content

Commit

Permalink
feat(lib/trie): atomic tracked merkle values
Browse files Browse the repository at this point in the history
  • Loading branch information
@qdm12
qdm12 committed Oct 6, 2022
1 parent f922506 commit e89b1cc
Show file tree
Hide file tree
Showing 2 changed files with 127 additions and 29 deletions.
65 changes: 52 additions & 13 deletions lib/trie/trie.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,11 @@ type Trie struct {
// pruner to detect with database keys (trie node Merkle values) can
// be deleted.
deletedMerkleValues map[string]struct{}
// newDeletedMerkleValues contains deleted Merkle values for operations
// in progress. Its goal is to prevent having the trie `deletedMerkleValues`
// in a bad state if one of the trie operation modified some (deep copied) trie
// node and then failed.
newDeletedMerkleValues map[string]struct{}
}

// NewEmptyTrie creates a trie with a nil root
Expand All @@ -36,10 +41,11 @@ func NewEmptyTrie() *Trie {
// NewTrie creates a trie with an existing root node
func NewTrie(root *Node) *Trie {
return &Trie{
root: root,
childTries: make(map[common.Hash]*Trie),
generation: 0, // Initially zero but increases after every snapshot.
deletedMerkleValues: make(map[string]struct{}),
root: root,
childTries: make(map[common.Hash]*Trie),
generation: 0, // Initially zero but increases after every snapshot.
deletedMerkleValues: make(map[string]struct{}),
newDeletedMerkleValues: make(map[string]struct{}),
}
}

Expand All @@ -54,28 +60,42 @@ func (t *Trie) Snapshot() (newTrie *Trie) {
rootCopySettings.CopyCached = true
for rootHash, childTrie := range t.childTries {
childTries[rootHash] = &Trie{
generation: childTrie.generation + 1,
root: childTrie.root.Copy(rootCopySettings),
deletedMerkleValues: make(map[string]struct{}),
generation: childTrie.generation + 1,
root: childTrie.root.Copy(rootCopySettings),
deletedMerkleValues: make(map[string]struct{}),
newDeletedMerkleValues: make(map[string]struct{}),
}
}

return &Trie{
generation: t.generation + 1,
root: t.root,
childTries: childTries,
deletedMerkleValues: make(map[string]struct{}),
generation: t.generation + 1,
root: t.root,
childTries: childTries,
deletedMerkleValues: make(map[string]struct{}),
newDeletedMerkleValues: make(map[string]struct{}),
}
}

// handleTrackedDeltas modifies the trie deleted merkle values set
// only if the error is nil. In all cases, it resets the
// `newDeletedMerkleValues` map.
func (t *Trie) handleTrackedDeltas(err error) {
if err == nil {
for merkleValue := range t.newDeletedMerkleValues {
t.deletedMerkleValues[merkleValue] = struct{}{}
}
}
t.newDeletedMerkleValues = make(map[string]struct{})
}

func (t *Trie) prepLeafForMutation(currentLeaf *Node,
copySettings node.CopySettings) (newLeaf *Node) {
if currentLeaf.Generation == t.generation {
// no need to deep copy and update generation
// of current leaf.
newLeaf = currentLeaf
} else {
newLeaf = updateGeneration(currentLeaf, t.generation, t.deletedMerkleValues, copySettings)
newLeaf = updateGeneration(currentLeaf, t.generation, t.newDeletedMerkleValues, copySettings)
}
newLeaf.SetDirty()
return newLeaf
Expand All @@ -88,7 +108,7 @@ func (t *Trie) prepBranchForMutation(currentBranch *Node,
// of current branch.
newBranch = currentBranch
} else {
newBranch = updateGeneration(currentBranch, t.generation, t.deletedMerkleValues, copySettings)
newBranch = updateGeneration(currentBranch, t.generation, t.newDeletedMerkleValues, copySettings)
}
newBranch.SetDirty()
return newBranch
Expand Down Expand Up @@ -322,6 +342,9 @@ func findNextKeyChild(children []*Node, startIndex byte,
// Put inserts a value into the trie at the
// key specified in little Endian format.
func (t *Trie) Put(keyLE, value []byte) {
defer func() {
t.handleTrackedDeltas(nil)
}()
nibblesKey := codec.KeyLEToNibbles(keyLE)
t.root, _, _ = t.insert(t.root, nibblesKey, value)
}
Expand Down Expand Up @@ -515,6 +538,10 @@ func (t *Trie) insertInBranch(parentBranch *Node, key, value []byte) (
// The keys are in hexadecimal little Endian encoding and the values
// are hexadecimal encoded.
func (t *Trie) LoadFromMap(data map[string]string) (err error) {
defer func() {
t.handleTrackedDeltas(err)
}()

for key, value := range data {
keyLEBytes, err := common.HexToBytes(key)
if err != nil {
Expand Down Expand Up @@ -681,6 +708,10 @@ func retrieveFromBranch(branch *Node, key []byte) (value []byte) {
// keys and a boolean indicating if all keys with the prefix were deleted
// within the limit.
func (t *Trie) ClearPrefixLimit(prefixLE []byte, limit uint32) (deleted uint32, allDeleted bool) {
defer func() {
t.handleTrackedDeltas(nil)
}()

if limit == 0 {
return 0, false
}
Expand Down Expand Up @@ -865,6 +896,10 @@ func (t *Trie) deleteNodesLimit(parent *Node, limit uint32) (
// ClearPrefix deletes all nodes in the trie for which the key contains the
// prefix given in little Endian format.
func (t *Trie) ClearPrefix(prefixLE []byte) {
defer func() {
t.handleTrackedDeltas(nil)
}()

if len(prefixLE) == 0 {
t.root = nil
return
Expand Down Expand Up @@ -950,6 +985,10 @@ func (t *Trie) clearPrefixAtNode(parent *Node, prefix []byte) (
// matching the key given in little Endian format.
// If no node is found at this key, nothing is deleted.
func (t *Trie) Delete(keyLE []byte) {
defer func() {
t.handleTrackedDeltas(nil)
}()

key := codec.KeyLEToNibbles(keyLE)
t.root, _, _ = t.deleteAtNode(t.root, key)
}
Expand Down
91 changes: 75 additions & 16 deletions lib/trie/trie_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,9 @@ import (

func Test_NewEmptyTrie(t *testing.T) {
expectedTrie := &Trie{
childTries: make(map[common.Hash]*Trie),
deletedMerkleValues: map[string]struct{}{},
childTries: make(map[common.Hash]*Trie),
deletedMerkleValues: map[string]struct{}{},
newDeletedMerkleValues: map[string]struct{}{},
}
trie := NewEmptyTrie()
assert.Equal(t, expectedTrie, trie)
Expand All @@ -35,8 +36,9 @@ func Test_NewTrie(t *testing.T) {
Key: []byte{0},
SubValue: []byte{17},
},
childTries: make(map[common.Hash]*Trie),
deletedMerkleValues: map[string]struct{}{},
childTries: make(map[common.Hash]*Trie),
deletedMerkleValues: map[string]struct{}{},
newDeletedMerkleValues: map[string]struct{}{},
}
trie := NewTrie(root)
assert.Equal(t, expectedTrie, trie)
Expand All @@ -55,42 +57,55 @@ func Test_Trie_Snapshot(t *testing.T) {
deletedMerkleValues: map[string]struct{}{
"a": {},
},
newDeletedMerkleValues: map[string]struct{}{
"a temp": {},
},
},
{2}: {
generation: 2,
root: &Node{Key: []byte{2}, SubValue: []byte{1}},
deletedMerkleValues: map[string]struct{}{
"b": {},
},
newDeletedMerkleValues: map[string]struct{}{
"b temp": {},
},
},
},
deletedMerkleValues: map[string]struct{}{
"a": {},
"b": {},
},
newDeletedMerkleValues: map[string]struct{}{
"a temp": {},
"b temp": {},
},
}

expectedTrie := &Trie{
generation: 9,
root: &Node{Key: []byte{8}, SubValue: []byte{1}},
childTries: map[common.Hash]*Trie{
{1}: {
generation: 2,
root: &Node{Key: []byte{1}, SubValue: []byte{1}},
deletedMerkleValues: map[string]struct{}{},
generation: 2,
root: &Node{Key: []byte{1}, SubValue: []byte{1}},
deletedMerkleValues: map[string]struct{}{},
newDeletedMerkleValues: map[string]struct{}{},
},
{2}: {
generation: 3,
root: &Node{Key: []byte{2}, SubValue: []byte{1}},
deletedMerkleValues: map[string]struct{}{},
generation: 3,
root: &Node{Key: []byte{2}, SubValue: []byte{1}},
deletedMerkleValues: map[string]struct{}{},
newDeletedMerkleValues: map[string]struct{}{},
},
},
deletedMerkleValues: map[string]struct{}{},
deletedMerkleValues: map[string]struct{}{},
newDeletedMerkleValues: map[string]struct{}{},
}

newTrie := trie.Snapshot()

assert.Equal(t, expectedTrie, newTrie)
assert.Equal(t, expectedTrie.childTries, newTrie.childTries)
}

func Test_Trie_updateGeneration(t *testing.T) {
Expand Down Expand Up @@ -1007,7 +1022,8 @@ func Test_Trie_Put(t *testing.T) {
}{
"trie with key and value": {
trie: Trie{
generation: 1,
generation: 1,
newDeletedMerkleValues: map[string]struct{}{},
root: &Node{
Key: []byte{1, 2, 0, 5},
SubValue: []byte{1},
Expand All @@ -1016,7 +1032,8 @@ func Test_Trie_Put(t *testing.T) {
key: []byte{0x12, 0x16},
value: []byte{2},
expectedTrie: Trie{
generation: 1,
generation: 1,
newDeletedMerkleValues: map[string]struct{}{},
root: &Node{
Key: []byte{1, 2},
Generation: 1,
Expand Down Expand Up @@ -1578,21 +1595,34 @@ func Test_Trie_LoadFromMap(t *testing.T) {
errWrapped error
errMessage string
}{
"nil data": {},
"nil data": {
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
},
"empty data": {
data: map[string]string{},
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
},
"bad key": {
data: map[string]string{
"0xa": "0x01",
},
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
errWrapped: hex.ErrLength,
errMessage: "cannot convert key hex to bytes: encoding/hex: odd length hex string: 0xa",
},
"bad value": {
data: map[string]string{
"0x01": "0xa",
},
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
errWrapped: hex.ErrLength,
errMessage: "cannot convert value hex to bytes: encoding/hex: odd length hex string: 0xa",
},
Expand All @@ -1603,6 +1633,7 @@ func Test_Trie_LoadFromMap(t *testing.T) {
"0x0130": "0x08",
},
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
root: &Node{
Key: []byte{00, 01},
SubValue: []byte{6},
Expand Down Expand Up @@ -1650,6 +1681,7 @@ func Test_Trie_LoadFromMap(t *testing.T) {
"0x0130": "0x08",
},
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
root: &Node{
Key: []byte{00, 01},
SubValue: []byte{6},
Expand Down Expand Up @@ -2138,7 +2170,11 @@ func Test_Trie_ClearPrefixLimit(t *testing.T) {
allDeleted bool
expectedTrie Trie
}{
"limit is zero": {},
"limit is zero": {
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
},
"clear prefix limit": {
trie: Trie{
root: &Node{
Expand All @@ -2158,6 +2194,9 @@ func Test_Trie_ClearPrefixLimit(t *testing.T) {
limit: 5,
deleted: 2,
allDeleted: true,
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
},
}

Expand Down Expand Up @@ -2900,15 +2939,24 @@ func Test_Trie_ClearPrefix(t *testing.T) {
trie: Trie{
root: &Node{SubValue: []byte{1}},
},
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
},
"empty prefix": {
trie: Trie{
root: &Node{SubValue: []byte{1}},
},
prefix: []byte{},
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
},
"empty trie": {
prefix: []byte{0x12},
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
},
"clear prefix": {
trie: Trie{
Expand Down Expand Up @@ -2940,6 +2988,7 @@ func Test_Trie_ClearPrefix(t *testing.T) {
SubValue: []byte{1},
Dirty: true,
},
newDeletedMerkleValues: map[string]struct{}{},
},
},
}
Expand Down Expand Up @@ -3271,14 +3320,23 @@ func Test_Trie_Delete(t *testing.T) {
trie: Trie{
root: &Node{SubValue: []byte{1}},
},
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
},
"empty key": {
trie: Trie{
root: &Node{SubValue: []byte{1}},
},
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
},
"empty trie": {
key: []byte{0x12},
expectedTrie: Trie{
newDeletedMerkleValues: map[string]struct{}{},
},
},
"delete branch node": {
trie: Trie{
Expand Down Expand Up @@ -3326,6 +3384,7 @@ func Test_Trie_Delete(t *testing.T) {
},
}),
},
newDeletedMerkleValues: map[string]struct{}{},
},
},
}
Expand Down

0 comments on commit e89b1cc

Please sign in to comment.