From 853aa68d926089d98b83344214bff308d9e73688 Mon Sep 17 00:00:00 2001
From: cec489 <173723251+cec489@users.noreply.github.com>
Date: Tue, 23 Jul 2024 19:30:07 +0000
Subject: [PATCH] Remove heap allocations when flushing the nodes or cached
 leaves

---
 .../utreexobackends/cachedleavesmap.go        | 14 ++++++++++
 .../internal/utreexobackends/nodesmap.go      | 14 ++++++++++
 blockchain/utreexoio.go                       | 28 +++++++++----------
 3 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/blockchain/internal/utreexobackends/cachedleavesmap.go b/blockchain/internal/utreexobackends/cachedleavesmap.go
index cc46f26d..f18111bd 100644
--- a/blockchain/internal/utreexobackends/cachedleavesmap.go
+++ b/blockchain/internal/utreexobackends/cachedleavesmap.go
@@ -133,6 +133,20 @@ func (ms *CachedLeavesMapSlice) DeleteMaps() {
 	}
 }
 
+// ClearMaps clears all maps
+//
+// This function is safe for concurrent access.
+func (ms *CachedLeavesMapSlice) ClearMaps() {
+	ms.mtx.Lock()
+	defer ms.mtx.Unlock()
+
+	for i := range ms.maps {
+		for key := range ms.maps[i] {
+			delete(ms.maps[i], key)
+		}
+	}
+}
+
 // ForEach loops through all the elements in the cachedleaves map slice and calls fn with the key-value pairs.
 //
 // This function is safe for concurrent access.
diff --git a/blockchain/internal/utreexobackends/nodesmap.go b/blockchain/internal/utreexobackends/nodesmap.go
index edf6693d..0c2a753c 100644
--- a/blockchain/internal/utreexobackends/nodesmap.go
+++ b/blockchain/internal/utreexobackends/nodesmap.go
@@ -170,6 +170,20 @@ func (ms *NodesMapSlice) DeleteMaps() {
 	}
 }
 
+// ClearMaps clears all maps
+//
+// This function is safe for concurrent access.
+func (ms *NodesMapSlice) ClearMaps() {
+	ms.mtx.Lock()
+	defer ms.mtx.Unlock()
+
+	for i := range ms.maps {
+		for key := range ms.maps[i] {
+			delete(ms.maps[i], key)
+		}
+	}
+}
+
 // ForEach loops through all the elements in the nodes map slice and calls fn with the key-value pairs.
 //
 // This function is safe for concurrent access.
diff --git a/blockchain/utreexoio.go b/blockchain/utreexoio.go
index 28b22d23..0c0a6fd9 100644
--- a/blockchain/utreexoio.go
+++ b/blockchain/utreexoio.go
@@ -16,6 +16,9 @@ import (
 // leafLength is the length of a seriailzed leaf.
 const leafLength = chainhash.HashSize + 1
 
+// buffer size for VLQ serialization.  Double the size needed to serialize 2^64
+const vlqBufSize = 22
+
 // serializeLeaf serializes the leaf to [leafLength]byte.
 func serializeLeaf(leaf utreexo.Leaf) [leafLength]byte {
 	var buf [leafLength]byte
@@ -68,12 +71,11 @@ func InitNodesBackEnd(datadir string, maxTotalMemoryUsage int64) (*NodesBackEnd,
 
 // dbPut serializes and puts the key value pair into the database.
 func (m *NodesBackEnd) dbPut(k uint64, v utreexo.Leaf) error {
-	size := serializeSizeVLQ(k)
-	buf := make([]byte, size)
-	putVLQ(buf, k)
+	var buf [vlqBufSize]byte
+	size := putVLQ(buf[:], k)
 
 	serialized := serializeLeaf(v)
-	return m.db.Put(buf[:], serialized[:], nil)
+	return m.db.Put(buf[:size], serialized[:], nil)
 }
 
 // dbGet fetches the value from the database and deserializes it and returns
@@ -98,10 +100,9 @@ func (m *NodesBackEnd) dbGet(k uint64) (utreexo.Leaf, bool) {
 
 // dbDel removes the key from the database.
 func (m *NodesBackEnd) dbDel(k uint64) error {
-	size := serializeSizeVLQ(k)
-	buf := make([]byte, size)
-	putVLQ(buf, k)
-	return m.db.Delete(buf, nil)
+	var buf [vlqBufSize]byte
+	size := putVLQ(buf[:], k)
+	return m.db.Delete(buf[:size], nil)
 }
 
 // Get returns the leaf from the underlying map.
@@ -276,7 +277,7 @@ func (m *NodesBackEnd) flush() {
 		}
 	})
 
-	m.cache.DeleteMaps()
+	m.cache.ClearMaps()
 }
 
 // Close flushes the cache and closes the underlying database.
@@ -298,10 +299,9 @@ type CachedLeavesBackEnd struct {
 
 // dbPut serializes and puts the key and the value into the database.
 func (m *CachedLeavesBackEnd) dbPut(k utreexo.Hash, v uint64) error {
-	size := serializeSizeVLQ(v)
-	buf := make([]byte, size)
-	putVLQ(buf, v)
-	return m.db.Put(k[:], buf, nil)
+	var buf [vlqBufSize]byte
+	size := putVLQ(buf[:], v)
+	return m.db.Put(k[:], buf[:size], nil)
 }
 
 // dbGet fetches and deserializes the value from the database.
@@ -411,7 +411,7 @@ func (m *CachedLeavesBackEnd) flush() {
 		}
 	})
 
-	m.cache.DeleteMaps()
+	m.cache.ClearMaps()
 }
 
 // Close flushes all the cached entries and then closes the underlying database.