From 65d2676ad24d66c43e0cdfc0b68cce749c774c8a Mon Sep 17 00:00:00 2001
From: thomassong <thomassong2012@gmail.com>
Date: Wed, 24 Jan 2024 15:14:25 +0800
Subject: [PATCH] fix(txn): discard empty transactions on CommitWith (#2031)

---
 batch_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 txn.go        |  2 ++
 2 files changed, 43 insertions(+)

diff --git a/batch_test.go b/batch_test.go
index 59b39ad06..189f5fbde 100644
--- a/batch_test.go
+++ b/batch_test.go
@@ -104,3 +104,44 @@ func TestBatchErrDeadlock(t *testing.T) {
 	require.Error(t, wb.Flush())
 	require.NoError(t, db.Close())
 }
+
+// This test ensures we don't end up in deadlock in case of empty writebatch.
+func TestEmptyWriteBatch(t *testing.T) {
+	t.Run("normal mode", func(t *testing.T) {
+		runBadgerTest(t, nil, func(t *testing.T, db *DB) {
+			wb := db.NewWriteBatch()
+			require.NoError(t, wb.Flush())
+			wb = db.NewWriteBatch()
+			require.NoError(t, wb.Flush())
+			wb = db.NewWriteBatch()
+			// Flush commits inner txn and sets a new one instead.
+			// Thus we need to save it to check if it was discarded.
+			txn := wb.txn
+			require.NoError(t, wb.Flush())
+			// check that flushed txn was discarded and marked as read.
+			require.True(t, txn.discarded)
+		})
+	})
+	t.Run("managed mode", func(t *testing.T) {
+		opt := getTestOptions("")
+		opt.managedTxns = true
+		runBadgerTest(t, &opt, func(t *testing.T, db *DB) {
+			t.Run("WriteBatchAt", func(t *testing.T) {
+				wb := db.NewWriteBatchAt(2)
+				require.NoError(t, wb.Flush())
+				wb = db.NewWriteBatchAt(208)
+				require.NoError(t, wb.Flush())
+				wb = db.NewWriteBatchAt(31)
+				require.NoError(t, wb.Flush())
+			})
+			t.Run("ManagedWriteBatch", func(t *testing.T) {
+				wb := db.NewManagedWriteBatch()
+				require.NoError(t, wb.Flush())
+				wb = db.NewManagedWriteBatch()
+				require.NoError(t, wb.Flush())
+				wb = db.NewManagedWriteBatch()
+				require.NoError(t, wb.Flush())
+			})
+		})
+	})
+}
diff --git a/txn.go b/txn.go
index bdee3c877..9fb7d9d4a 100644
--- a/txn.go
+++ b/txn.go
@@ -622,6 +622,8 @@ func (txn *Txn) CommitWith(cb func(error)) {
 		// callback might be acquiring the same locks. Instead run the callback
 		// from another goroutine.
 		go runTxnCallback(&txnCb{user: cb, err: nil})
+		// Discard the transaction so that the read is marked done.
+		txn.Discard()
 		return
 	}