Merge pull request #148 from kcalvinalvin/2024-03-07-add-caching-to-u…

…treexo-backends

blockchain, indexers: Add cache option to bridgenodes

blockchain/indexers/flatutreexoproofindex.go

@@ -1243,12 +1243,14 @@ func loadFlatFileState(dataDir, name string) (*FlatFileState, error) {
 }
 
 // NewFlatUtreexoProofIndex returns a new instance of an indexer that is used to create a flat utreexo proof index.
+// The passed in maxMemoryUsage should be in bytes and it determines how much memory the proof index will use up.
+// A maxMemoryUsage of 0 will keep all the elements on disk and a negative maxMemoryUsage will keep all the elements in memory.
 //
 // It implements the Indexer interface which plugs into the IndexManager that in
 // turn is used by the blockchain package.  This allows the index to be
 // seamlessly maintained along with the chain.
-func NewFlatUtreexoProofIndex(dataDir string, pruned bool, chainParams *chaincfg.Params,
-	proofGenInterVal *int32) (*FlatUtreexoProofIndex, error) {
+func NewFlatUtreexoProofIndex(pruned bool, chainParams *chaincfg.Params,
+	proofGenInterVal *int32, maxMemoryUsage int64, dataDir string) (*FlatUtreexoProofIndex, error) {
 
 	// If the proofGenInterVal argument is nil, use the default value.
 	var intervalToUse int32
@@ -1269,9 +1271,8 @@ func NewFlatUtreexoProofIndex(dataDir string, pruned bool, chainParams *chaincfg
 	uState, err := InitUtreexoState(&UtreexoConfig{
 		DataDir: dataDir,
 		Name:    flatUtreexoProofIndexType,
-		// Default to ram for now.
-		Params: chainParams,
-	})
+		Params:  chainParams,
+	}, maxMemoryUsage)
 	if err != nil {
 		return nil, err
 	}

blockchain/indexers/indexers_test.go

@@ -75,12 +75,12 @@ func initIndexes(interval int32, dbPath string, db *database.DB, params *chaincf
 
 	proofGenInterval := new(int32)
 	*proofGenInterval = interval
-	flatUtreexoProofIndex, err := NewFlatUtreexoProofIndex(dbPath, false, params, proofGenInterval)
+	flatUtreexoProofIndex, err := NewFlatUtreexoProofIndex(false, params, proofGenInterval, 50*1024*1024, dbPath)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	utreexoProofIndex, err := NewUtreexoProofIndex(*db, false, dbPath, params)
+	utreexoProofIndex, err := NewUtreexoProofIndex(*db, false, 50*1024*1024, params, dbPath)
 	if err != nil {
 		return nil, nil, err
 	}

blockchain/indexers/utreexobackend.go

@@ -7,6 +7,7 @@ package indexers
 import (
 	"bytes"
 	"encoding/binary"
+	"fmt"
 	"os"
 	"path/filepath"
 
@@ -63,11 +64,13 @@ func utreexoBasePath(cfg *UtreexoConfig) string {
 
 // InitUtreexoState returns an initialized utreexo state. If there isn't an
 // existing state on disk, it creates one and returns it.
-func InitUtreexoState(cfg *UtreexoConfig) (*UtreexoState, error) {
+// maxMemoryUsage of 0 will keep every element on disk. A negaive maxMemoryUsage will
+// load every element to the memory.
+func InitUtreexoState(cfg *UtreexoConfig, maxMemoryUsage int64) (*UtreexoState, error) {
 	basePath := utreexoBasePath(cfg)
 	log.Infof("Initializing Utreexo state from '%s'", basePath)
 	defer log.Info("Utreexo state loaded")
-	return initUtreexoState(cfg, basePath)
+	return initUtreexoState(cfg, maxMemoryUsage, basePath)
 }
 
 // deleteUtreexoState removes the utreexo state directory and all the contents
@@ -311,18 +314,24 @@ func deserializeUndoBlock(serialized []byte) (uint64, []uint64, []utreexo.Hash,
 	return numAdds, targets, delHashes, nil
 }
 
-// initUtreexoState creates a new utreexo state and returns it.
-func initUtreexoState(cfg *UtreexoConfig, basePath string) (*UtreexoState, error) {
+// initUtreexoState creates a new utreexo state and returns it. maxMemoryUsage of 0 will keep
+// every element on disk and a negative maxMemoryUsage will load all the elemnts to memory.
+func initUtreexoState(cfg *UtreexoConfig, maxMemoryUsage int64, basePath string) (*UtreexoState, error) {
 	p := utreexo.NewMapPollard(true)
 
+	// 60% of the memory for the nodes map, 40% for the cache leaves map.
+	// TODO Totally arbitrary, it there's something better than change it to that.
+	maxNodesMem := maxMemoryUsage * 6 / 10
+	maxCachedLeavesMem := maxMemoryUsage - maxNodesMem
+
 	nodesPath := filepath.Join(basePath, nodesDBDirName)
-	nodesDB, err := blockchain.InitNodesBackEnd(nodesPath)
+	nodesDB, err := blockchain.InitNodesBackEnd(nodesPath, maxNodesMem)
 	if err != nil {
 		return nil, err
 	}
 
 	cachedLeavesPath := filepath.Join(basePath, cachedLeavesDBDirName)
-	cachedLeavesDB, err := blockchain.InitCachedLeavesBackEnd(cachedLeavesPath)
+	cachedLeavesDB, err := blockchain.InitCachedLeavesBackEnd(cachedLeavesPath, maxCachedLeavesMem)
 	if err != nil {
 		return nil, err
 	}
@@ -341,20 +350,76 @@ func initUtreexoState(cfg *UtreexoConfig, basePath string) (*UtreexoState, error
 		p.NumLeaves = binary.LittleEndian.Uint64(buf[:])
 	}
 
-	p.Nodes = nodesDB
-	p.CachedLeaves = cachedLeavesDB
-	closeDB := func() error {
-		err := nodesDB.Close()
+	var closeDB func() error
+	if maxMemoryUsage >= 0 {
+		p.Nodes = nodesDB
+		p.CachedLeaves = cachedLeavesDB
+		closeDB = func() error {
+			err := nodesDB.Close()
+			if err != nil {
+				return err
+			}
+
+			err = cachedLeavesDB.Close()
+			if err != nil {
+				return err
+			}
+
+			return nil
+		}
+	} else {
+		log.Infof("loading the utreexo state from disk...")
+		err = nodesDB.ForEach(func(k uint64, v utreexo.Leaf) error {
+			p.Nodes.Put(k, v)
+			return nil
+		})
 		if err != nil {
-			return err
+			return nil, err
 		}
 
-		err = cachedLeavesDB.Close()
+		err = cachedLeavesDB.ForEach(func(k utreexo.Hash, v uint64) error {
+			p.CachedLeaves.Put(k, v)
+			return nil
+		})
 		if err != nil {
-			return err
+			return nil, err
 		}
 
-		return nil
+		log.Infof("Finished loading the utreexo state from disk.")
+
+		closeDB = func() error {
+			log.Infof("Flushing the utreexo state to disk. May take a while...")
+
+			p.Nodes.ForEach(func(k uint64, v utreexo.Leaf) error {
+				nodesDB.Put(k, v)
+				return nil
+			})
+
+			p.CachedLeaves.ForEach(func(k utreexo.Hash, v uint64) error {
+				cachedLeavesDB.Put(k, v)
+				return nil
+			})
+
+			// We want to try to close both of the DBs before returning because of an error.
+			errStr := ""
+			err := nodesDB.Close()
+			if err != nil {
+				errStr += fmt.Sprintf("Error while closing nodes db. %v", err.Error())
+			}
+			err = cachedLeavesDB.Close()
+			if err != nil {
+				errStr += fmt.Sprintf("Error while closing cached leaves db. %v", err.Error())
+			}
+
+			// If the err string isn't "", then return the error here.
+			if errStr != "" {
+				return fmt.Errorf(errStr)
+			}
+
+			log.Infof("Finished flushing the utreexo state to disk.")
+
+			return nil
+		}
 	}
 
 	uState := &UtreexoState{

blockchain/indexers/utreexoproofindex.go

@@ -559,12 +559,17 @@ func (idx *UtreexoProofIndex) PruneBlock(dbTx database.Tx, blockHash *chainhash.
 	return nil
 }
 
-// NewUtreexoProofIndex returns a new instance of an indexer that is used to create a
+// NewUtreexoProofIndex returns a new instance of an indexer that is used to create a utreexo
+// proof index using the database passed in. The passed in maxMemoryUsage should be in bytes and
+// it determines how much memory the proof index will use up. A maxMemoryUsage of 0 will keep
+// all the elements on disk and a negative maxMemoryUsage will keep all the elements in memory.
 //
 // It implements the Indexer interface which plugs into the IndexManager that in
 // turn is used by the blockchain package.  This allows the index to be
 // seamlessly maintained along with the chain.
-func NewUtreexoProofIndex(db database.DB, pruned bool, dataDir string, chainParams *chaincfg.Params) (*UtreexoProofIndex, error) {
+func NewUtreexoProofIndex(db database.DB, pruned bool, maxMemoryUsage int64,
+	chainParams *chaincfg.Params, dataDir string) (*UtreexoProofIndex, error) {
+
 	idx := &UtreexoProofIndex{
 		db:          db,
 		chainParams: chainParams,
@@ -575,7 +580,7 @@ func NewUtreexoProofIndex(db database.DB, pruned bool, dataDir string, chainPara
 		DataDir: dataDir,
 		Name:    db.Type(),
 		Params:  chainParams,
-	})
+	}, maxMemoryUsage)
 	if err != nil {
 		return nil, err
 	}