chore: custom backup draft

Author: fracasula

internal/cache/badger/badger.go

@@ -1,58 +1,63 @@
 package badger
 
 import (
+	"bytes"
 	"compress/gzip"
+	"context"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
-	"os"
-	"path"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/dgraph-io/badger/v4"
 	"github.com/dgraph-io/badger/v4/options"
+	"github.com/dgraph-io/badger/v4/pb"
+	"github.com/dgraph-io/ristretto/v2/z"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/rudderlabs/rudder-go-kit/bytesize"
 	"github.com/rudderlabs/rudder-go-kit/config"
 	"github.com/rudderlabs/rudder-go-kit/logger"
+	obskit "github.com/rudderlabs/rudder-observability-kit/go/labels"
 )
 
 var ErrSnapshotInProgress = errors.New("snapshotting already in progress")
 
+type hasher interface {
+	GetKeysByHashRange(keys []string) (
+		map[uint32][]string, // itemsByHashRange
+		error,
+	)
+	GetKeysByHashRangeWithIndexes(keys []string) (
+		map[uint32][]string, // itemsByHashRange
+		map[string]int, // indexes
+		error,
+	)
+}
+
 type Cache struct {
+	hasher           hasher
 	cache            *badger.DB
+	conf             *config.Config
+	logger           logger.Logger
 	compress         bool
 	discardRatio     float64
 	snapshotSince    uint64
 	snapshotting     bool
 	snapshottingLock sync.Mutex
+	debugMode        bool
 }
 
-func Factory(conf *config.Config, log logger.Logger) func(hashRange uint32) (*Cache, error) {
-	return func(hashRange uint32) (*Cache, error) {
-		badgerPath := conf.GetString("BadgerDB.Dedup.Path", "/tmp/badger")
-		badgerPath = path.Join(badgerPath, fmt.Sprintf("%d", hashRange))
-		if err := os.MkdirAll(badgerPath, 0o755); err != nil {
-			return nil, fmt.Errorf("failed to create badger directory for hash range %d: %w", hashRange, err)
-		}
-		if conf.GetBool("BadgerDB.TestMode", false) {
-			// TODO use this to enable usage of String() string and Len() methods
-		}
-		badgerCache, err := New(badgerPath, conf, log)
-		if err != nil {
-			return nil, fmt.Errorf("failed to create cache factory: %w", err)
-		}
-		return badgerCache, nil
-	}
-}
-
-func New(path string, conf *config.Config, log logger.Logger) (*Cache, error) {
+func New(h hasher, conf *config.Config, log logger.Logger) (*Cache, error) {
+	path := conf.GetString("BadgerDB.Dedup.Path", "/tmp/badger")
 	opts := badger.DefaultOptions(path).
 		WithCompression(options.None).
-		WithNumGoroutines(1).
 		WithNumVersionsToKeep(1).
+		WithNumGoroutines(conf.GetInt("BadgerDB.Dedup.NumGoroutines", 128/3)).
 		WithBloomFalsePositive(conf.GetFloat64("BadgerDB.Dedup.BloomFalsePositive", 0.000001)).
 		WithIndexCacheSize(conf.GetInt64Var(16*bytesize.MB, 1, "BadgerDB.Dedup.indexCacheSize", "BadgerDB.indexCacheSize")).
 		WithValueLogFileSize(conf.GetInt64Var(1*bytesize.MB, 1, "BadgerDB.Dedup.valueLogFileSize", "BadgerDB.valueLogFileSize")).
@@ -65,7 +70,7 @@ func New(path string, conf *config.Config, log logger.Logger) (*Cache, error) {
 		WithBaseLevelSize(conf.GetInt64Var(5*bytesize.MB, 1, "BadgerDB.Dedup.baseLevelSize", "BadgerDB.baseLevelSize")).
 		WithLevelSizeMultiplier(conf.GetIntVar(10, 1, "BadgerDB.Dedup.levelSizeMultiplier", "BadgerDB.levelSizeMultiplier")).
 		WithMaxLevels(conf.GetIntVar(7, 1, "BadgerDB.Dedup.maxLevels", "BadgerDB.maxLevels")).
-		WithNumCompactors(conf.GetIntVar(2, 1, "BadgerDB.Dedup.numCompactors", "BadgerDB.numCompactors")).
+		WithNumCompactors(conf.GetIntVar(4, 1, "BadgerDB.Dedup.numCompactors", "BadgerDB.numCompactors")).
 		WithValueThreshold(conf.GetInt64Var(10*bytesize.B, 1, "BadgerDB.Dedup.valueThreshold", "BadgerDB.valueThreshold")).
 		WithSyncWrites(conf.GetBoolVar(false, "BadgerDB.Dedup.syncWrites", "BadgerDB.syncWrites")).
 		WithBlockCacheSize(conf.GetInt64Var(0, 1, "BadgerDB.Dedup.blockCacheSize", "BadgerDB.blockCacheSize")).
@@ -84,28 +89,39 @@ func New(path string, conf *config.Config, log logger.Logger) (*Cache, error) {
 		return nil, err
 	}
 	return &Cache{
+		hasher:       h,
 		cache:        db,
+		conf:         conf,
+		logger:       log,
 		compress:     compress,
 		discardRatio: conf.GetFloat64("BadgerDB.Dedup.DiscardRatio", 0.7),
+		debugMode:    conf.GetBool("BadgerDB.DebugMode", false),
 	}, nil
 }
 
 // Get returns the values associated with the keys and an error if the operation failed
 func (c *Cache) Get(keys []string) ([]bool, error) {
+	itemsByHashRange, indexes, err := c.hasher.GetKeysByHashRangeWithIndexes(keys)
+	if err != nil {
+		return nil, fmt.Errorf("cache get keys: %w", err)
+	}
+
 	results := make([]bool, len(keys))
 
-	err := c.cache.View(func(txn *badger.Txn) error {
-		for i, key := range keys {
-			_, err := txn.Get([]byte(key))
-			if err != nil {
-				if errors.Is(err, badger.ErrKeyNotFound) {
-					results[i] = false
-					continue
+	err = c.cache.View(func(txn *badger.Txn) error {
+		for hashRange, keys := range itemsByHashRange {
+			for _, key := range keys {
+				_, err := txn.Get(c.getKey(key, hashRange))
+				if err != nil {
+					if errors.Is(err, badger.ErrKeyNotFound) {
+						results[indexes[key]] = false
+						continue
+					}
+					return fmt.Errorf("failed to get key %s: %w", key, err)
 				}
-				return fmt.Errorf("failed to get key %s: %w", key, err)
+				// Key exists, value is true
+				results[indexes[key]] = true
 			}
-			// Key exists, value is true
-			results[i] = true
 		}
 		return nil
 	})
@@ -118,27 +134,39 @@ func (c *Cache) Get(keys []string) ([]bool, error) {
 
 // Put adds or updates elements inside the cache with the specified TTL and returns an error if the operation failed
 func (c *Cache) Put(keys []string, ttl time.Duration) error {
-	err := c.cache.Update(func(txn *badger.Txn) error {
-		for _, key := range keys {
-			entry := badger.NewEntry([]byte(key), []byte{})
-			if ttl > 0 {
-				entry = entry.WithTTL(ttl)
-			}
-			if err := txn.SetEntry(entry); err != nil {
-				return fmt.Errorf("failed to put key %s: %w", key, err)
+	itemsByHashRange, err := c.hasher.GetKeysByHashRange(keys)
+	if err != nil {
+		return fmt.Errorf("cache put keys: %w", err)
+	}
+
+	err = c.cache.Update(func(txn *badger.Txn) error {
+		for hashRange, keys := range itemsByHashRange {
+			for _, key := range keys {
+				entry := badger.NewEntry(c.getKey(key, hashRange), []byte{})
+				if ttl > 0 {
+					entry = entry.WithTTL(ttl)
+				}
+				if err := txn.SetEntry(entry); err != nil {
+					return fmt.Errorf("failed to put key %s: %w", key, err)
+				}
 			}
 		}
 		return nil
 	})
 	if err != nil {
 		return err
 	}
+
 	return nil
 }
 
 // Len returns the number of elements in the cache
 // WARNING: this must be used in tests only TODO protect this with a testMode=false by default
 func (c *Cache) Len() int {
+	if !c.debugMode {
+		panic("Len() is only available in debug mode")
+	}
+
 	count := 0
 	err := c.cache.View(func(txn *badger.Txn) error {
 		opts := badger.DefaultIteratorOptions
@@ -160,6 +188,10 @@ func (c *Cache) Len() int {
 // String returns a string representation of the cache
 // WARNING: this must be used in tests only TODO protect this with a testMode=false by default
 func (c *Cache) String() string {
+	if !c.debugMode {
+		panic("String() is only available in debug mode")
+	}
+
 	sb := strings.Builder{}
 	sb.WriteString("{")
 
@@ -191,7 +223,7 @@ func (c *Cache) String() string {
 }
 
 // CreateSnapshot writes the cache contents to the provided writer
-func (c *Cache) CreateSnapshot(w io.Writer) (uint64, error) {
+func (c *Cache) CreateSnapshot(ctx context.Context, w map[uint32]io.Writer) (uint64, error) {
 	c.snapshottingLock.Lock()
 	if c.snapshotting {
 		c.snapshottingLock.Unlock()
@@ -204,20 +236,142 @@ func (c *Cache) CreateSnapshot(w io.Writer) (uint64, error) {
 
 	// TODO gzip should be configurable, maybe we want to use another algorithm with a different compression level
 	if c.compress {
-		w = gzip.NewWriter(w)
-		defer func() { _ = w.(*gzip.Writer).Close() }()
+		for hashRange, writer := range w {
+			w[hashRange] = gzip.NewWriter(writer)
+		}
+		defer func() {
+			for _, writer := range w {
+				err := writer.(*gzip.Writer).Close()
+				if err != nil {
+					c.logger.Errorn("failed to close gzip writer", obskit.Error(err))
+				}
+			}
+		}()
 	}
 
-	newSince, err := c.cache.Backup(w, since)
+	hashRangesMap := make(map[uint32]struct{})
+	for hashRange := range w {
+		hashRangesMap[hashRange] = struct{}{}
+	}
+
+	var (
+		maxVersion     uint64
+		keysToDelete   = make([][]byte, 0)
+		keysToDeleteMu = sync.Mutex{}
+	)
+	stream := c.cache.NewStream()
+	stream.NumGo = c.conf.GetInt("BadgerDB.Dedup.Snapshots.NumGoroutines", 10)
+	stream.Prefix = []byte("hr")
+	stream.SinceTs = since
+	stream.ChooseKey = func(item *badger.Item) (ok bool) {
+		hasExpired := item.ExpiresAt() > 0 && item.ExpiresAt() <= uint64(time.Now().Unix())
+		if !hasExpired {
+			parts := bytes.Split(item.Key(), []byte(":"))
+			hashRange, _ := strconv.ParseUint(string(parts[0][2:]), 10, 32)
+			_, ok = hashRangesMap[uint32(hashRange)]
+		}
+		if hasExpired || !ok {
+			keysToDeleteMu.Lock()
+			keysToDelete = append(keysToDelete, item.Key())
+			keysToDeleteMu.Unlock()
+			return false
+		}
+		return true
+	}
+	stream.Send = func(buf *z.Buffer) error {
+		list, err := badger.BufferToKVList(buf)
+		if err != nil {
+			return err
+		}
+
+		// Group KV pairs by hash range
+		kvsByHashRange := make(map[uint32][]*pb.KV)
+
+		for _, kv := range list.Kv {
+			if maxVersion < kv.Version {
+				maxVersion = kv.Version
+			}
+			if !kv.StreamDone {
+				// Extract the hash range from the key.
+				// Key format is "hr<hashRange>:<actualKey>".
+				parts := bytes.Split(kv.Key, []byte(":"))
+				if len(parts) < 2 {
+					c.logger.Warnn("Skipping malformed key", logger.NewStringField("key", string(kv.Key)))
+					continue // Skip malformed keys
+				}
+
+				hashRangeStr := string(parts[0][2:]) // Remove "hr" prefix
+				hashRange, err := strconv.ParseUint(hashRangeStr, 10, 32)
+				if err != nil {
+					c.logger.Warnn("Skipping key with invalid hash range", logger.NewStringField("key", string(kv.Key)))
+					continue // Skip keys with invalid hash range
+				}
+
+				hashRange32 := uint32(hashRange)
+
+				// Only include if we have a writer for this hash range
+				if _, exists := w[hashRange32]; exists {
+					kvsByHashRange[hashRange32] = append(kvsByHashRange[hashRange32], kv)
+				}
+			}
+		}
+
+		// Write to appropriate writers by hash range
+		for hashRange, kvs := range kvsByHashRange {
+			writer := w[hashRange]
+
+			// Create a new KVList for this hash range
+			rangeList := &pb.KVList{Kv: kvs}
+
+			if err := writeTo(rangeList, writer); err != nil {
+				return fmt.Errorf("failed to write to hash range %d: %w", hashRange, err)
+			}
+		}
+
+		return nil
+	}
+
+	err := stream.Orchestrate(ctx)
 	if err != nil {
 		return 0, fmt.Errorf("failed to create snapshot: %w", err)
 	}
 
 	c.snapshottingLock.Lock()
-	c.snapshotSince = newSince
+	c.snapshotSince = maxVersion
 	c.snapshotting = false
 	c.snapshottingLock.Unlock()
 
+	if len(keysToDelete) > 0 {
+		batchSize := c.conf.GetInt("BadgerDB.Dedup.Snapshots.DeleteBatchSize", 1000)
+		keysToDeleteBatch := make([][]byte, 0, batchSize)
+		deleteKeys := func() error {
+			err := c.cache.Update(func(txn *badger.Txn) error {
+				for _, key := range keysToDeleteBatch {
+					return txn.Delete(key)
+				}
+				return nil
+			})
+			if err != nil {
+				return fmt.Errorf("failed to delete keys: %w", err)
+			}
+			keysToDeleteBatch = make([][]byte, 0, batchSize)
+			return nil
+		}
+		for _, key := range keysToDelete {
+			keysToDeleteBatch = append(keysToDeleteBatch, key)
+			if len(keysToDeleteBatch) == batchSize {
+				if err := deleteKeys(); err != nil {
+					return 0, fmt.Errorf("failed to delete keys: %w", err)
+				}
+			}
+		}
+		if len(keysToDeleteBatch) > 0 {
+			if err := deleteKeys(); err != nil {
+				return 0, fmt.Errorf("failed to delete keys: %w", err)
+			}
+		}
+	}
+
 	return since, nil
 }
 
@@ -247,6 +401,22 @@ func (c *Cache) Close() error {
 	return c.cache.Close()
 }
 
+func (c *Cache) getKey(key string, hashRange uint32) []byte {
+	return []byte("hr" + strconv.Itoa(int(hashRange)) + ":" + key)
+}
+
+func writeTo(list *pb.KVList, w io.Writer) error {
+	if err := binary.Write(w, binary.LittleEndian, uint64(proto.Size(list))); err != nil {
+		return err
+	}
+	buf, err := proto.Marshal(list)
+	if err != nil {
+		return err
+	}
+	_, err = w.Write(buf)
+	return err
+}
+
 type loggerForBadger struct {
 	logger.Logger
 }