chore(lib/trie): refactor and accelerate existing end to end tests (#…

…2211)

.gitignore

@@ -20,6 +20,7 @@
 *.wasm
 
 test_data
+trie_putandget_failed_test_data_*
 tmp
 
 tests/utils/config*

dot/state/base_test.go

@@ -4,7 +4,6 @@
 package state
 
 import (
-	"bytes"
 	"testing"
 
 	"github.com/ChainSafe/gossamer/lib/common"
@@ -18,14 +17,13 @@ func TestTrie_StoreAndLoadFromDB(t *testing.T) {
 	db := NewInMemoryDB(t)
 	tt := trie.NewEmptyTrie()
 
-	rt := trie.GenerateRandomTests(t, 1000)
-	for _, test := range rt {
-		tt.Put(test.Key(), test.Value())
+	generator := newGenerator()
+	const size = 500
+	kv := generateKeyValues(t, generator, size)
 
-		val := tt.Get(test.Key())
-		if !bytes.Equal(val, test.Value()) {
-			t.Errorf("Fail to get key %x with value %x: got %x", test.Key(), test.Value(), val)
-		}
+	for keyString, value := range kv {
+		key := []byte(keyString)
+		tt.Put(key, value)
 	}
 
 	err := tt.Store(db)

dot/state/helpers_test.go

@@ -0,0 +1,75 @@
+// Copyright 2022 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package state
+
+import (
+	"math/rand"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// newGenerator creates a new PRNG seeded with the
+// unix nanoseconds value of the current time.
+func newGenerator() (prng *rand.Rand) {
+	seed := time.Now().UnixNano()
+	source := rand.NewSource(seed)
+	return rand.New(source)
+}
+
+func generateKeyValues(tb testing.TB, generator *rand.Rand, size int) (kv map[string][]byte) {
+	tb.Helper()
+
+	kv = make(map[string][]byte, size)
+
+	const maxKeySize, maxValueSize = 510, 128
+	for i := 0; i < size; i++ {
+		populateKeyValueMap(tb, kv, generator, maxKeySize, maxValueSize)
+	}
+
+	return kv
+}
+
+func populateKeyValueMap(tb testing.TB, kv map[string][]byte,
+	generator *rand.Rand, maxKeySize, maxValueSize int) {
+	tb.Helper()
+
+	for {
+		const minKeySize = 2
+		key := generateRandBytesMinMax(tb, minKeySize, maxKeySize, generator)
+
+		keyString := string(key)
+
+		_, keyExists := kv[keyString]
+
+		if keyExists && key[1] != byte(0) {
+			continue
+		}
+
+		const minValueSize = 2
+		value := generateRandBytesMinMax(tb, minValueSize, maxValueSize, generator)
+
+		kv[keyString] = value
+
+		break
+	}
+}
+
+func generateRandBytesMinMax(tb testing.TB, minSize, maxSize int,
+	generator *rand.Rand) (b []byte) {
+	tb.Helper()
+	size := minSize +
+		generator.Intn(maxSize-minSize)
+	return generateRandBytes(tb, size, generator)
+}
+
+func generateRandBytes(tb testing.TB, size int,
+	generator *rand.Rand) (b []byte) {
+	tb.Helper()
+	b = make([]byte, size)
+	_, err := generator.Read(b)
+	require.NoError(tb, err)
+	return b
+}

lib/trie/helpers_test.go

@@ -5,6 +5,11 @@ package trie
 
 import (
 	"errors"
+	"math/rand"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
 )
 
 type writeCall struct {
@@ -14,3 +19,72 @@ type writeCall struct {
 }
 
 var errTest = errors.New("test error")
+
+type Test struct {
+	key   []byte
+	value []byte
+	op    int
+}
+
+// newGenerator creates a new PRNG seeded with the
+// unix nanoseconds value of the current time.
+func newGenerator() (prng *rand.Rand) {
+	seed := time.Now().UnixNano()
+	source := rand.NewSource(seed)
+	return rand.New(source)
+}
+
+func generateKeyValues(tb testing.TB, generator *rand.Rand, size int) (kv map[string][]byte) {
+	tb.Helper()
+
+	kv = make(map[string][]byte, size)
+
+	const maxKeySize, maxValueSize = 510, 128
+	for i := 0; i < size; i++ {
+		populateKeyValueMap(tb, kv, generator, maxKeySize, maxValueSize)
+	}
+
+	return kv
+}
+
+func populateKeyValueMap(tb testing.TB, kv map[string][]byte,
+	generator *rand.Rand, maxKeySize, maxValueSize int) {
+	tb.Helper()
+
+	for {
+		const minKeySize = 2
+		key := generateRandBytesMinMax(tb, minKeySize, maxKeySize, generator)
+
+		keyString := string(key)
+
+		_, keyExists := kv[keyString]
+
+		if keyExists && key[1] != byte(0) {
+			continue
+		}
+
+		const minValueSize = 2
+		value := generateRandBytesMinMax(tb, minValueSize, maxValueSize, generator)
+
+		kv[keyString] = value
+
+		break
+	}
+}
+
+func generateRandBytesMinMax(tb testing.TB, minSize, maxSize int,
+	generator *rand.Rand) (b []byte) {
+	tb.Helper()
+	size := minSize +
+		generator.Intn(maxSize-minSize)
+	return generateRandBytes(tb, size, generator)
+}
+
+func generateRandBytes(tb testing.TB, size int,
+	generator *rand.Rand) (b []byte) {
+	tb.Helper()
+	b = make([]byte, size)
+	_, err := generator.Read(b)
+	require.NoError(tb, err)
+	return b
+}

lib/trie/mem_test.go

@@ -0,0 +1,114 @@
+// Copyright 2022 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package trie
+
+import (
+	"math"
+	"runtime"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_Trie_MemoryUsage(t *testing.T) {
+	// Set skip to false to run the test.
+	// This test should be run on its own since it interacts
+	// with the Go garbage collector.
+	const skip = true
+	if skip {
+		t.SkipNow()
+	}
+
+	triesMap := map[string]*Trie{
+		"first": NewEmptyTrie(),
+	}
+
+	generator := newGenerator()
+	const size = 10000
+	kv := generateKeyValues(t, generator, size)
+
+	// Populate a first branch branching out
+	// from the root on the 'left'
+	populateTrieAtPrefix(triesMap["first"], []byte{0, 1}, kv)
+
+	// Check heap memory usage - it should be X
+	halfFilledTrieHeap := getHeapUsage()
+
+	// Populate a second branch branching out
+	// from the root on the 'right'
+	populateTrieAtPrefix(triesMap["first"], []byte{0, 2}, kv)
+
+	// Check heap memory usage - it should be 2X
+	filledTrieHeap := getHeapUsage()
+	ratio := getApproximateRatio(halfFilledTrieHeap, filledTrieHeap)
+	assert.Greater(t, ratio, 1.5)
+	assert.Less(t, ratio, 2.1)
+
+	// Snapshot the trie
+	triesMap["second"] = triesMap["first"].Snapshot()
+
+	// Modify all the leaves from the first branch in the new trie
+	mutateTrieLeavesAtPrefix(triesMap["second"], []byte{0, 1}, kv)
+
+	// Check heap memory usage - it should be 3X
+	halfMutatedTrieHeap := getHeapUsage()
+	ratio = getApproximateRatio(halfFilledTrieHeap, halfMutatedTrieHeap)
+	assert.Greater(t, ratio, 2.0)
+	assert.Less(t, ratio, 3.1)
+
+	// Remove the older trie from our reference
+	delete(triesMap, "first")
+
+	// Check heap memory usage - it should be 2X
+	prunedTrieHeap := getHeapUsage()
+	ratio = getApproximateRatio(halfFilledTrieHeap, prunedTrieHeap)
+	assert.Greater(t, ratio, 1.5)
+	assert.Less(t, ratio, 2.1)
+
+	// Dummy calls - has to be after prunedTrieHeap for
+	// GC to keep them
+	_, ok := triesMap["first"]
+	require.False(t, ok)
+	_, ok = kv["dummy"]
+	require.False(t, ok)
+}
+
+func getApproximateRatio(old, new uint64) (ratio float64) {
+	ratio = float64(new) / float64(old)
+	ratio = math.Round(ratio*100) / 100
+	return ratio
+}
+
+func getHeapUsage() (heapAlloc uint64) {
+	runtime.GC()
+	var memStats runtime.MemStats
+	runtime.ReadMemStats(&memStats)
+	return memStats.HeapAlloc
+}
+
+func populateTrieAtPrefix(trie *Trie,
+	prefix []byte, kv map[string][]byte) {
+	for keyString, value := range kv {
+		key := append(prefix, []byte(keyString)...)
+
+		trie.Put(key, value)
+	}
+}
+
+func mutateTrieLeavesAtPrefix(trie *Trie,
+	prefix []byte, originalKV map[string][]byte) {
+	for keyString, value := range originalKV {
+		key := append(prefix, []byte(keyString)...)
+
+		// Reverse value byte slice
+		newValue := make([]byte, len(value))
+		copy(newValue, value)
+		for i, j := 0, len(newValue)-1; i < j; i, j = i+1, j-1 {
+			newValue[i], newValue[j] = newValue[j], newValue[i]
+		}
+
+		trie.Put(key, value)
+	}
+}

lib/trie/proof_test.go

@@ -21,14 +21,17 @@ func TestProofGeneration(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	expectedValue := rand32Bytes()
+	const size = 32
+	generator := newGenerator()
+
+	expectedValue := generateRandBytes(t, size, generator)
 
 	trie := NewEmptyTrie()
-	trie.Put([]byte("cat"), rand32Bytes())
-	trie.Put([]byte("catapulta"), rand32Bytes())
+	trie.Put([]byte("cat"), generateRandBytes(t, size, generator))
+	trie.Put([]byte("catapulta"), generateRandBytes(t, size, generator))
 	trie.Put([]byte("catapora"), expectedValue)
-	trie.Put([]byte("dog"), rand32Bytes())
-	trie.Put([]byte("doguinho"), rand32Bytes())
+	trie.Put([]byte("dog"), generateRandBytes(t, size, generator))
+	trie.Put([]byte("doguinho"), generateRandBytes(t, size, generator))
 
 	err = trie.Store(memdb)
 	require.NoError(t, err)