From 9683021e745b01b35e64ce5e7a99311955481ac1 Mon Sep 17 00:00:00 2001
From: "siddarth.msv" <82795818+Sidddddarth@users.noreply.github.com>
Date: Thu, 17 Oct 2024 11:42:06 +0530
Subject: [PATCH] fix: dedup - use writeBatch to put keys in badger (#5204)

---
 services/dedup/badger/badger.go | 20 ++++++--------------
 1 file changed, 6 insertions(+), 14 deletions(-)

diff --git a/services/dedup/badger/badger.go b/services/dedup/badger/badger.go
index 768d5bdfd8..92e0bfe5e9 100644
--- a/services/dedup/badger/badger.go
+++ b/services/dedup/badger/badger.go
@@ -101,25 +101,17 @@ func (d *BadgerDB) Get(key string) (int64, bool, error) {
 
 func (d *BadgerDB) Set(kvs []types.KeyValue) error {
 	defer d.stats.NewTaggedStat("dedup_commit_duration_seconds", stats.TimerType, stats.Tags{"mode": "badger"}).RecordDuration()()
-
-	txn := d.badgerDB.NewTransaction(true)
-	for _, message := range kvs {
+	wb := d.badgerDB.NewWriteBatch()
+	defer wb.Cancel()
+	for i := range kvs {
+		message := kvs[i]
 		value := strconv.FormatInt(message.Value, 10)
 		e := badger.NewEntry([]byte(message.Key), []byte(value)).WithTTL(d.window.Load())
-		err := txn.SetEntry(e)
-		if err == badger.ErrTxnTooBig {
-			if err = txn.Commit(); err != nil {
-				return err
-			}
-			txn = d.badgerDB.NewTransaction(true)
-			if err = txn.SetEntry(e); err != nil {
-				return err
-			}
-		} else if err != nil {
+		if err := wb.SetEntry(e); err != nil {
 			return err
 		}
 	}
-	return txn.Commit()
+	return wb.Flush()
 }
 
 func (d *BadgerDB) Close() {