Merkle db intermediate node key compression

x/merkledb/db.go

@@ -1188,21 +1188,28 @@ func (db *merkleDB) getNode(key path, hasValue bool) (*node, error) {
 // The returned []byte is taken from [bufferPool] and
 // should be returned to it when the caller is done with it.
 func addPrefixToKey(bufferPool *sync.Pool, prefix []byte, key []byte) []byte {
-	prefixedKey := bufferPool.Get().([]byte)
 	prefixLen := len(prefix)
 	keyLen := prefixLen + len(key)
-	if cap(prefixedKey) >= keyLen {
+	prefixedKey := getBufferFromPool(bufferPool, keyLen)
+	copy(prefixedKey, prefix)
+	copy(prefixedKey[prefixLen:], key)
+	return prefixedKey
+}
+
+// Returns a []byte from [bufferPool] with length exactly [size].
+// The []byte is not guaranteed to be zeroed.
+func getBufferFromPool(bufferPool *sync.Pool, size int) []byte {
+	buffer := bufferPool.Get().([]byte)
+	if cap(buffer) >= size {
 		// The [] byte we got from the pool is big enough to hold the prefixed key
-		prefixedKey = prefixedKey[:keyLen]
+		buffer = buffer[:size]
 	} else {
 		// The []byte from the pool wasn't big enough.
 		// Put it back and allocate a new, bigger one
-		bufferPool.Put(prefixedKey)
-		prefixedKey = make([]byte, keyLen)
+		bufferPool.Put(buffer)
+		buffer = make([]byte, size)
 	}
-	copy(prefixedKey, prefix)
-	copy(prefixedKey[prefixLen:], key)
-	return prefixedKey
+	return buffer
 }
 
 // cacheEntrySize returns a rough approximation of the memory consumed by storing the path and node

x/merkledb/intermediate_node_db.go

@@ -89,13 +89,13 @@ func (db *intermediateNodeDB) onEviction(key path, n *node) error {
 }
 
 func (db *intermediateNodeDB) addToBatch(b database.Batch, key path, n *node) error {
-	prefixedKey := addPrefixToKey(db.bufferPool, intermediateNodePrefix, key.Bytes())
-	defer db.bufferPool.Put(prefixedKey)
+	dbKey := db.constructDBKey(key)
+	defer db.bufferPool.Put(dbKey)
 	db.metrics.DatabaseNodeWrite()
 	if n == nil {
-		return b.Delete(prefixedKey)
+		return b.Delete(dbKey)
 	}
-	return b.Put(prefixedKey, n.bytes())
+	return b.Put(dbKey, n.bytes())
 }
 
 func (db *intermediateNodeDB) Get(key path) (*node, error) {
@@ -108,17 +108,44 @@ func (db *intermediateNodeDB) Get(key path) (*node, error) {
 	}
 	db.metrics.IntermediateNodeCacheMiss()
 
-	prefixedKey := addPrefixToKey(db.bufferPool, intermediateNodePrefix, key.Bytes())
+	dbKey := db.constructDBKey(key)
 	db.metrics.DatabaseNodeRead()
-	nodeBytes, err := db.baseDB.Get(prefixedKey)
+	nodeBytes, err := db.baseDB.Get(dbKey)
 	if err != nil {
 		return nil, err
 	}
-	db.bufferPool.Put(prefixedKey)
+	db.bufferPool.Put(dbKey)
 
 	return parseNode(key, nodeBytes)
 }
 
+// constructDBKey returns a key that can be used in [db.baseDB].
+// We need to be able to differentiate between two paths of equal
+// byte length but different token length so we add padding to differentiate.
+// Additionally, we add a prefix indicating it is part of the intermediateNodeDB.
+func (db *intermediateNodeDB) constructDBKey(key path) []byte {
+	compressedKey := key.Serialize()
+
+	// add one additional byte to store padding when the path
+	// has a length that fits into a whole number of bytes
+	remainder := compressedKey.NibbleLength % 2
+	keyLen := len(compressedKey.Value)
+	if remainder == 0 {
+		keyLen++
+	}
+	dbKey := getBufferFromPool(db.bufferPool, keyLen)
+	defer db.bufferPool.Put(dbKey)
+
+	copy(dbKey, compressedKey.Value)
+	if remainder == 0 {
+		dbKey[keyLen-1] = 0b1000_0000
+	} else {
+		dbKey[keyLen-1] |= 0b0000_1000
+	}
+
+	return addPrefixToKey(db.bufferPool, intermediateNodePrefix, dbKey)
+}
+
 func (db *intermediateNodeDB) Put(key path, n *node) error {
 	return db.nodeCache.Put(key, n)
 }

x/merkledb/intermediate_node_db_test.go

@@ -20,7 +20,7 @@ import (
 // * Deleting a key-node pair from the database
 // * Evicting elements from the cache
 // * Flushing the cache
-func TestIntermediateNodeDB(t *testing.T) {
+func Test_IntermediateNodeDB(t *testing.T) {
 	require := require.New(t)
 
 	// use exact multiple of node size so require.Equal(1, db.nodeCache.fifo.Len()) is correct later
@@ -70,7 +70,7 @@ func TestIntermediateNodeDB(t *testing.T) {
 	added := 0
 	for {
 		key := newPath([]byte{byte(added)})
-		node := newNode(nil, EmptyPath)
+		node := newNode(nil, key)
 		node.setValue(maybe.Some([]byte{byte(added)}))
 		newExpectedSize := expectedSize + cacheEntrySize(key, node)
 		if newExpectedSize > cacheSize {
@@ -90,7 +90,7 @@ func TestIntermediateNodeDB(t *testing.T) {
 	// of all but 2 elements. 2 elements remain rather than 1 element because of
 	// the added key prefix increasing the size tracked by the batch.
 	key := newPath([]byte{byte(added)})
-	node := newNode(nil, EmptyPath)
+	node := newNode(nil, key)
 	node.setValue(maybe.Some([]byte{byte(added)}))
 	require.NoError(db.Put(key, node))
 
@@ -125,3 +125,72 @@ func TestIntermediateNodeDB(t *testing.T) {
 	require.NoError(it.Error())
 	require.Equal(added+1, count)
 }
+
+func FuzzIntermediateNodeDBConstructDBKey(f *testing.F) {
+	cacheSize := 200
+	evictionBatchSize := cacheSize
+	baseDB := memdb.New()
+	db := newIntermediateNodeDB(
+		baseDB,
+		&sync.Pool{
+			New: func() interface{} { return make([]byte, 0) },
+		},
+		&mockMetrics{},
+		cacheSize,
+		evictionBatchSize,
+	)
+
+	f.Fuzz(func(
+		t *testing.T,
+		key []byte,
+	) {
+		require := require.New(t)
+
+		p := newPath(key)
+		constructedKey := db.constructDBKey(p)
+		baseLength := len(p)/2 + len(intermediateNodePrefix)
+		require.Equal(intermediateNodePrefix, constructedKey[:len(intermediateNodePrefix)])
+		if len(p)%2 == 0 {
+			// when even, there is an extra padding byte
+			require.Len(constructedKey, baseLength+1)
+			require.Equal(append(p.Serialize().Value, 0b1000_0000), constructedKey[len(intermediateNodePrefix):])
+		} else {
+			require.Len(constructedKey, baseLength)
+			require.Equal(p.Append(0b0000_1000).Serialize().Value, constructedKey[len(intermediateNodePrefix):])
+		}
+	})
+}
+
+func Test_IntermediateNodeDB_ConstructDBKey_DirtyBuffer(t *testing.T) {
+	require := require.New(t)
+	cacheSize := 200
+	evictionBatchSize := cacheSize
+	baseDB := memdb.New()
+	db := newIntermediateNodeDB(
+		baseDB,
+		&sync.Pool{
+			New: func() interface{} { return make([]byte, 0) },
+		},
+		&mockMetrics{},
+		cacheSize,
+		evictionBatchSize,
+	)
+
+	db.bufferPool.Put([]byte{0xFF, 0xFF, 0xFF})
+	constructedKey := db.constructDBKey(newPath([]byte{}))
+	require.Len(constructedKey, 2)
+	require.Equal(intermediateNodePrefix, constructedKey[:len(intermediateNodePrefix)])
+	require.Equal(byte(0b1000_0000), constructedKey[len(constructedKey)-1])
+
+	db.bufferPool = &sync.Pool{
+		New: func() interface{} {
+			return make([]byte, 0, defaultBufferLength)
+		},
+	}
+	db.bufferPool.Put([]byte{0xFF, 0xFF, 0xFF})
+	p := path(" ")
+	constructedKey = db.constructDBKey(p)
+	require.Len(constructedKey, 2)
+	require.Equal(intermediateNodePrefix, constructedKey[:len(intermediateNodePrefix)])
+	require.Equal(p.Append(0b0000_1000).Serialize().Value, constructedKey[len(intermediateNodePrefix):])
+}