Implement BadgerGC with controller runtime

Signed-off-by: leigh capili <[email protected]>

Author: stealthybox

internal/database/badger_gc.go

@@ -0,0 +1,93 @@
+/*
+Copyright 2025 The Flux authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package database
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/dgraph-io/badger/v3"
+	"github.com/go-logr/logr"
+	ctrl "sigs.k8s.io/controller-runtime"
+)
+
+// BadgerGarbageCollector implements controller runtime's Runnable
+type BadgerGarbageCollector struct {
+	// DiscardRatio must be a float between 0.0 and 1.0, inclusive
+	// See badger.DB.RunValueLogGC for more info
+	DiscardRatio float64
+	Interval     time.Duration
+
+	name string
+	db   *badger.DB
+	log  logr.Logger
+}
+
+// NewBadgerGarbageCollector creates and returns a new BadgerGarbageCollector
+func NewBadgerGarbageCollector(name string, db *badger.DB, interval time.Duration, discardRatio float64) *BadgerGarbageCollector {
+	return &BadgerGarbageCollector{
+		DiscardRatio: discardRatio,
+		Interval:     interval,
+
+		name: name,
+		db:   db,
+	}
+}
+
+// Start repeatedly runs the BadgerDB garbage collector with a delay inbetween
+// runs.
+//
+// Start blocks until the context is cancelled. The database is expected to
+// already be open and not be closed while this context is active.
+//
+// ctx should be a logr.Logger context.
+func (gc *BadgerGarbageCollector) Start(ctx context.Context) error {
+	gc.log = ctrl.LoggerFrom(ctx).WithName(gc.name)
+
+	gc.log.Info("Starting Badger GC")
+	timer := time.NewTimer(gc.Interval)
+	for {
+		select {
+		case <-timer.C:
+			gc.discardValueLogFiles()
+			timer.Reset(gc.Interval)
+		case <-ctx.Done():
+			timer.Stop()
+			gc.log.Info("Stopped Badger GC")
+			return nil
+		}
+	}
+}
+
+// upper bound for loop
+const maxDiscards = 1000
+
+func (gc *BadgerGarbageCollector) discardValueLogFiles() {
+	for c := 0; c < maxDiscards; c++ {
+		err := gc.db.RunValueLogGC(gc.DiscardRatio)
+		if errors.Is(err, badger.ErrNoRewrite) {
+			// there is no more garbage to discard
+			gc.log.Info("Ran Badger GC", "discarded_vlogs", c)
+			return
+		}
+		if err != nil {
+			gc.log.Error(err, "Badger GC Error", "discarded_vlogs", c)
+			return
+		}
+	}
+	gc.log.Info("Ran Badger GC for maximum discards", "discarded_vlogs", maxDiscards)
+}

internal/database/badger_gc_test.go

@@ -0,0 +1,76 @@
+/*
+Copyright 2020 The Flux authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package database
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/dgraph-io/badger/v3"
+	"github.com/go-logr/logr"
+	"github.com/go-logr/logr/testr"
+)
+
+func TestBadgerGarbageCollectorDoesStop(t *testing.T) {
+	badger, db := createBadgerDatabaseForGC(t)
+	ctx, cancel := context.WithCancel(
+		logr.NewContext(context.Background(), testr.New(t)))
+
+	stop := make(chan struct{})
+	go func() {
+		gc := NewBadgerGarbageCollector("test-badger-gc", badger, 500*time.Millisecond, 0.01)
+		gc.Start(ctx)
+		stop <- struct{}{}
+	}()
+
+	time.Sleep(time.Second)
+
+	tags := []string{"latest", "v0.0.1", "v0.0.2"}
+	fatalIfError(t, db.SetTags(testRepo, tags))
+	_, err := db.Tags(testRepo)
+	fatalIfError(t, err)
+	t.Log("wrote tags successfully")
+
+	time.Sleep(time.Second)
+
+	cancel()
+	t.Log("waiting for GC stop")
+	select {
+	case <-time.NewTimer(5 * time.Second).C:
+		t.Fatalf("GC did not stop")
+	case <-stop:
+		t.Log("GC Stopped")
+	}
+}
+
+func createBadgerDatabaseForGC(t *testing.T) (*badger.DB, *BadgerDatabase) {
+	t.Helper()
+	dir, err := os.MkdirTemp(os.TempDir(), "badger")
+	if err != nil {
+		t.Fatal(err)
+	}
+	db, err := badger.Open(badger.DefaultOptions(dir))
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		db.Close()
+		os.RemoveAll(dir)
+	})
+	return db, NewBadgerDatabase(db)
+}

main.go

@@ -20,6 +20,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"time"
 
 	"github.com/dgraph-io/badger/v3"
 	flag "github.com/spf13/pflag"
@@ -132,7 +133,9 @@ func main() {
 		os.Exit(1)
 	}
 	defer badgerDB.Close()
+
 	db := database.NewBadgerDatabase(badgerDB)
+	badgerGC := database.NewBadgerGarbageCollector("badger-gc", badgerDB, 1*time.Minute, 0.7)
 
 	watchNamespace := ""
 	if !watchOptions.AllNamespaces {
@@ -205,6 +208,8 @@ func main() {
 		os.Exit(1)
 	}
 
+	mgr.Add(badgerGC)
+
 	probes.SetupChecks(mgr, setupLog)
 
 	var eventRecorder *events.Recorder