-
Notifications
You must be signed in to change notification settings - Fork 112
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
1 changed file
with
111 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,111 @@ | ||
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) | ||
} | ||
} |