[BREAKING] opt(compactions): Improve compaction performance

badger/cmd/bank.go

@@ -357,7 +357,7 @@ func runTest(cmd *cobra.Command, args []string) error {
 	// Open DB
 	opts := badger.DefaultOptions(sstDir).
 		WithValueDir(vlogDir).
-		WithMaxTableSize(4 << 20). // Force more compactions.
+		WithBaseTableSize(4 << 20). // Force more compactions.
 		WithNumLevelZeroTables(2).
 		WithNumMemtables(2).
 		// Do not GC any versions, because we need them for the disect..

badger/cmd/write_bench.go

@@ -66,7 +66,7 @@ var (
 	vlogMaxEntries   uint32
 	loadBloomsOnOpen bool
 	detectConflicts  bool
-	compression      bool
+	zstdComp         bool
 	showDir          bool
 	ttlDuration      string
 	showKeysCount    bool
@@ -113,8 +113,8 @@ func init() {
 		"Load Bloom filter on DB open.")
 	writeBenchCmd.Flags().BoolVar(&detectConflicts, "conficts", false,
 		"If true, it badger will detect the conflicts")
-	writeBenchCmd.Flags().BoolVar(&compression, "compression", true,
-		"If true, badger will use ZSTD mode")
+	writeBenchCmd.Flags().BoolVar(&zstdComp, "zstd", false,
+		"If true, badger will use ZSTD mode. Otherwise, use default.")
 	writeBenchCmd.Flags().BoolVar(&showDir, "show-dir", false,
 		"If true, the report will include the directory contents")
 	writeBenchCmd.Flags().StringVar(&dropAllPeriod, "dropall", "0s",
@@ -260,12 +260,6 @@ func writeSorted(db *badger.DB, num uint64) error {
 }
 
 func writeBench(cmd *cobra.Command, args []string) error {
-	var cmode options.CompressionType
-	if compression {
-		cmode = options.ZSTD
-	} else {
-		cmode = options.None
-	}
 	opt := badger.DefaultOptions(sstDir).
 		WithValueDir(vlogDir).
 		WithSyncWrites(syncWrites).
@@ -277,8 +271,10 @@ func writeBench(cmd *cobra.Command, args []string) error {
 		WithValueLogMaxEntries(vlogMaxEntries).
 		WithEncryptionKey([]byte(encryptionKey)).
 		WithDetectConflicts(detectConflicts).
-		WithCompression(cmode).
 		WithLoggingLevel(badger.INFO)
+	if zstdComp {
+		opt = opt.WithCompression(options.ZSTD)
+	}
 
 	if !showLogs {
 		opt = opt.WithLogger(nil)
@@ -314,6 +310,7 @@ func writeBench(cmd *cobra.Command, args []string) error {
 	}
 
 	c.SignalAndWait()
+	fmt.Printf(db.LevelsToString())
 	return err
 }
 
@@ -354,11 +351,13 @@ func reportStats(c *z.Closer, db *badger.DB) {
 	t := time.NewTicker(time.Second)
 	defer t.Stop()
 
+	var count int
 	for {
 		select {
 		case <-c.HasBeenClosed():
 			return
 		case <-t.C:
+			count++
 			if showKeysCount {
 				showKeysStats(db)
 			}
@@ -392,8 +391,13 @@ func reportStats(c *z.Closer, db *badger.DB) {
 			bytesRate := sz / uint64(dur.Seconds())
 			entriesRate := entries / uint64(dur.Seconds())
 			fmt.Printf("[WRITE] Time elapsed: %s, bytes written: %s, speed: %s/sec, "+
-				"entries written: %d, speed: %d/sec, gcSuccess: %d\n", y.FixedDuration(time.Since(startTime)),
-				humanize.Bytes(sz), humanize.Bytes(bytesRate), entries, entriesRate, gcSuccess)
+				"entries written: %d, speed: %d/sec, Memory: %s\n",
+				y.FixedDuration(time.Since(startTime)),
+				humanize.Bytes(sz), humanize.Bytes(bytesRate), entries, entriesRate,
+				humanize.IBytes(uint64(z.NumAllocBytes())))
+			if count%10 == 0 {
+				fmt.Printf(db.LevelsToString())
+			}
 		}
 	}
 }

batch.go

@@ -82,14 +82,15 @@ func (wb *WriteBatch) Cancel() {
 	wb.txn.Discard()
 }
 
-// The caller of this callback must hold the lock.
 func (wb *WriteBatch) callback(err error) {
 	// sync.WaitGroup is thread-safe, so it doesn't need to be run inside wb.Lock.
 	defer wb.throttle.Done(err)
 	if err == nil {
 		return
 	}
 
+	wb.Lock()
+	defer wb.Unlock()
 	if wb.err != nil {
 		return
 	}

compaction.go

@@ -33,6 +33,10 @@ type keyRange struct {
 	inf   bool
 }
 
+func (r keyRange) isEmpty() bool {
+	return len(r.left) == 0 && len(r.right) == 0 && !r.inf
+}
+
 var infRange = keyRange{inf: true}
 
 func (r keyRange) String() string {
@@ -45,7 +49,26 @@ func (r keyRange) equals(dst keyRange) bool {
 		r.inf == dst.inf
 }
 
+func (r *keyRange) extend(kr keyRange) {
+	if r.isEmpty() {
+		*r = kr
+	}
+	if len(r.left) == 0 || y.CompareKeys(kr.left, r.left) < 0 {
+		r.left = kr.left
+	}
+	if len(r.right) == 0 || y.CompareKeys(kr.right, r.right) > 0 {
+		r.right = kr.right
+	}
+	if kr.inf {
+		r.inf = true
+	}
+}
+
 func (r keyRange) overlapsWith(dst keyRange) bool {
+	// Empty keyRange always overlaps.
+	if r.isEmpty() {
+		return true
+	}
 	if r.inf || dst.inf {
 		return true
 	}
@@ -127,6 +150,7 @@ func (lcs *levelCompactStatus) remove(dst keyRange) bool {
 type compactStatus struct {
 	sync.RWMutex
 	levels []*levelCompactStatus
+	tables map[uint64]struct{}
 }
 
 func (cs *compactStatus) overlapsWith(level int, this keyRange) bool {
@@ -151,11 +175,10 @@ func (cs *compactStatus) compareAndAdd(_ thisAndNextLevelRLocked, cd compactDef)
 	cs.Lock()
 	defer cs.Unlock()
 
-	level := cd.thisLevel.level
-
-	y.AssertTruef(level < len(cs.levels)-1, "Got level %d. Max levels: %d", level, len(cs.levels))
-	thisLevel := cs.levels[level]
-	nextLevel := cs.levels[level+1]
+	tl := cd.thisLevel.level
+	y.AssertTruef(tl < len(cs.levels)-1, "Got level %d. Max levels: %d", tl, len(cs.levels))
+	thisLevel := cs.levels[cd.thisLevel.level]
+	nextLevel := cs.levels[cd.nextLevel.level]
 
 	if thisLevel.overlapsWith(cd.thisRange) {
 		return false
@@ -171,31 +194,41 @@ func (cs *compactStatus) compareAndAdd(_ thisAndNextLevelRLocked, cd compactDef)
 	thisLevel.ranges = append(thisLevel.ranges, cd.thisRange)
 	nextLevel.ranges = append(nextLevel.ranges, cd.nextRange)
 	thisLevel.delSize += cd.thisSize
+	for _, t := range append(cd.top, cd.bot...) {
+		cs.tables[t.ID()] = struct{}{}
+	}
 	return true
 }
 
 func (cs *compactStatus) delete(cd compactDef) {
 	cs.Lock()
 	defer cs.Unlock()
 
-	level := cd.thisLevel.level
-	y.AssertTruef(level < len(cs.levels)-1, "Got level %d. Max levels: %d", level, len(cs.levels))
+	tl := cd.thisLevel.level
+	y.AssertTruef(tl < len(cs.levels)-1, "Got level %d. Max levels: %d", tl, len(cs.levels))
 
-	thisLevel := cs.levels[level]
-	nextLevel := cs.levels[level+1]
+	thisLevel := cs.levels[cd.thisLevel.level]
+	nextLevel := cs.levels[cd.nextLevel.level]
 
 	thisLevel.delSize -= cd.thisSize
 	found := thisLevel.remove(cd.thisRange)
-	found = nextLevel.remove(cd.nextRange) && found
+	if !cd.nextRange.isEmpty() {
+		found = nextLevel.remove(cd.nextRange) && found
+	}
 
 	if !found {
 		this := cd.thisRange
 		next := cd.nextRange
-		fmt.Printf("Looking for: [%q, %q, %v] in this level.\n", this.left, this.right, this.inf)
+		fmt.Printf("Looking for: %s in this level %d.\n", this, tl)
 		fmt.Printf("This Level:\n%s\n", thisLevel.debug())
 		fmt.Println()
-		fmt.Printf("Looking for: [%q, %q, %v] in next level.\n", next.left, next.right, next.inf)
+		fmt.Printf("Looking for: %s in next level %d.\n", next, cd.nextLevel.level)
 		fmt.Printf("Next Level:\n%s\n", nextLevel.debug())
 		log.Fatal("keyRange not found")
 	}
+	for _, t := range append(cd.top, cd.bot...) {
+		_, ok := cs.tables[t.ID()]
+		y.AssertTrue(ok)
+		delete(cs.tables, t.ID())
+	}
 }