From b4de33889e9f818ec74dbec2261511142e75a9c0 Mon Sep 17 00:00:00 2001
From: thomassong <thomassong2012@gmail.com>
Date: Thu, 25 Jan 2024 14:41:03 +0800
Subject: [PATCH] optimize allocations by using pre allocated buffer for
 priorities (#2006)

---
 db.go     |  2 +-
 levels.go | 17 ++++++++++++++---
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/db.go b/db.go
index d1f0813c8..3b9428acb 100644
--- a/db.go
+++ b/db.go
@@ -1334,7 +1334,7 @@ func (db *DB) Flatten(workers int) error {
 			}
 		}
 		if len(levels) <= 1 {
-			prios := db.lc.pickCompactLevels()
+			prios := db.lc.pickCompactLevels(nil)
 			if len(prios) == 0 || prios[0].score <= 1.0 {
 				db.opt.Infof("All tables consolidated into one level. Flattening done.\n")
 				return nil
diff --git a/levels.go b/levels.go
index c253a186c..49b11ad4f 100644
--- a/levels.go
+++ b/levels.go
@@ -397,8 +397,13 @@ func (s *levelsController) runCompactor(id int, lc *y.Closer) {
 		}
 		return false
 	}
+
+	var priosBuffer []compactionPriority
 	runOnce := func() bool {
-		prios := s.pickCompactLevels()
+		prios := s.pickCompactLevels(priosBuffer)
+		defer func() {
+			priosBuffer = prios
+		}()
 		if id == 0 {
 			// Worker ID zero prefers to compact L0 always.
 			prios = moveL0toFront(prios)
@@ -464,7 +469,7 @@ type compactionPriority struct {
 
 // pickCompactLevel determines which level to compact.
 // Based on: https://github.com/facebook/rocksdb/wiki/Leveled-Compaction
-func (s *levelsController) pickCompactLevels() (prios []compactionPriority) {
+func (s *levelsController) pickCompactLevels(priosBuffer []compactionPriority) (prios []compactionPriority) {
 	// This function must use identical criteria for guaranteeing compaction's progress that
 	// addLevel0Table uses.
 
@@ -477,6 +482,12 @@ func (s *levelsController) pickCompactLevels() (prios []compactionPriority) {
 		prios = append(prios, pri)
 	}
 
+	// Grow buffer to fit all levels.
+	if cap(priosBuffer) < len(s.levels) {
+		priosBuffer = make([]compactionPriority, 0, len(s.levels))
+	}
+	prios = priosBuffer[:0]
+
 	addPriority(0, float64(s.levels[0].numTables())/float64(s.kv.opt.NumLevelZeroTables))
 	for i := 1; i < len(s.levels); i++ {
 		// Don't consider those tables that are already being compacted right now.
@@ -1151,7 +1162,7 @@ func (s *levelsController) addLevel0Table(t *table.Table) error {
 			}
 			time.Sleep(10 * time.Millisecond)
 			if i%100 == 0 {
-				prios := s.pickCompactLevels()
+				prios := s.pickCompactLevels(nil)
 				s.elog.Printf("Waiting to add level 0 table. Compaction priorities: %+v\n", prios)
 				i = 0
 			}