From f3fe49ea5d3d9c4e9065cf504214c9f820e21ac7 Mon Sep 17 00:00:00 2001
From: Roman <ackhtariev@gmail.com>
Date: Thu, 17 Feb 2022 18:53:46 -0800
Subject: [PATCH] force GC, no cache during migration, auto heap profile

---
 mutable_tree.go | 57 ++++++++++++++++++++++++++++++++++++++++++++++++-
 nodedb.go       | 18 +++++++++++-----
 2 files changed, 69 insertions(+), 6 deletions(-)

diff --git a/mutable_tree.go b/mutable_tree.go
index 0c14af666..0dfe68e9e 100644
--- a/mutable_tree.go
+++ b/mutable_tree.go
@@ -4,8 +4,12 @@ import (
 	"bytes"
 	"crypto/sha256"
 	"fmt"
+	"os"
+	"runtime"
+	"runtime/pprof"
 	"sort"
 	"sync"
+	"time"
 
 	"github.com/pkg/errors"
 
@@ -534,6 +538,9 @@ func (tree *MutableTree) enableFastStorageAndCommitIfNotEnabled() (bool, error)
 		fastItr.Close()
 	}
 
+	// Force garbage collection before we proceed to enabling fast storage.
+	runtime.GC()
+
 	if err := tree.enableFastStorageAndCommit(); err != nil {
 		tree.ndb.storageVersion = defaultStorageVersionValue
 		return false, err
@@ -558,10 +565,58 @@ func (tree *MutableTree) enableFastStorageAndCommit() error {
 		}
 	}()
 
+	done := make(chan struct{})
+	defer func() {
+		done <- struct{}{}
+		close(done)
+	}()
+
+	go func ()  {
+		timer := time.NewTimer(time.Second)
+		defer func ()  {
+			if !timer.Stop() {
+				<-timer.C
+			}
+		}()
+
+		var m runtime.MemStats
+
+		hasTakenHeapProfile := false
+
+		for {
+			// Sample the current memory usage
+			runtime.ReadMemStats(&m)
+
+			if m.Alloc > 4 * 1024 * 1024 * 1024 {
+				// If we are using more than 4GB of memory, we should trigger garbage collection
+				// to free up some memory.
+				runtime.GC()
+				return
+			}
+
+			if !hasTakenHeapProfile && m.Alloc > 8 * 1024 * 1024 * 1024 {
+				// If we are using more than 8GB of memory, we should write a pprof sample
+				time := time.Now()
+				heapProfilePath := "/tmp/heap_profile" + time.String() + ".pprof"
+				heapFile, _ := os.Create(heapProfilePath)
+				pprof.WriteHeapProfile(heapFile)
+				heapFile.Close()
+				hasTakenHeapProfile = true
+			}
+
+			select {
+			case <-timer.C:
+				timer.Reset(time.Second)
+			case <-done:
+				return
+			}
+		}
+	}()
+
 	itr := NewIterator(nil, nil, true, tree.ImmutableTree)
 	defer itr.Close()
 	for ; itr.Valid(); itr.Next() {
-		if err = tree.ndb.SaveFastNode(NewFastNode(itr.Key(), itr.Value(), tree.version)); err != nil {
+		if err = tree.ndb.SaveFastNodeNoCache(NewFastNode(itr.Key(), itr.Value(), tree.version)); err != nil {
 			return err
 		}
 	}
diff --git a/nodedb.go b/nodedb.go
index c1843a7b4..7bee811c5 100644
--- a/nodedb.go
+++ b/nodedb.go
@@ -215,11 +215,18 @@ func (ndb *nodeDB) SaveNode(node *Node) {
 	ndb.cacheNode(node)
 }
 
-// SaveNode saves a FastNode to disk.
+// SaveNode saves a FastNode to disk and add to cache.
 func (ndb *nodeDB) SaveFastNode(node *FastNode) error {
 	ndb.mtx.Lock()
 	defer ndb.mtx.Unlock()
-	return ndb.saveFastNodeUnlocked(node)
+	return ndb.saveFastNodeUnlocked(node, true)
+}
+
+// SaveNode saves a FastNode to disk without adding to cache.
+func (ndb *nodeDB) SaveFastNodeNoCache(node *FastNode) error {
+	ndb.mtx.Lock()
+	defer ndb.mtx.Unlock()
+	return ndb.saveFastNodeUnlocked(node, false)
 }
 
 // setFastStorageVersionToBatch sets storage version to fast where the version is
@@ -274,7 +281,7 @@ func (ndb *nodeDB) shouldForceFastStorageUpgrade() bool {
 }
 
 // SaveNode saves a FastNode to disk.
-func (ndb *nodeDB) saveFastNodeUnlocked(node *FastNode) error {
+func (ndb *nodeDB) saveFastNodeUnlocked(node *FastNode, shouldAddToCache bool) error {
 	if node.key == nil {
 		return fmt.Errorf("FastNode cannot have a nil value for key")
 	}
@@ -290,8 +297,9 @@ func (ndb *nodeDB) saveFastNodeUnlocked(node *FastNode) error {
 	if err := ndb.batch.Set(ndb.fastNodeKey(node.key), buf.Bytes()); err != nil {
 		return fmt.Errorf("error while writing key/val to nodedb batch. Err: %w", err)
 	}
-	debug("BATCH SAVE %X %p\n", node.key, node)
-	ndb.cacheFastNode(node)
+	if shouldAddToCache {
+		ndb.cacheFastNode(node)
+	}
 	return nil
 }