From a2d675ce9588aff8a7c8de29dce8bca569b7d5c4 Mon Sep 17 00:00:00 2001
From: fatelei <fatelei@gmail.com>
Date: Fri, 4 Aug 2023 14:31:32 +0800
Subject: [PATCH] fix: Panic when closing DB while other goroutines doing reads

Signed-off-by: fatelei <fatelei@gmail.com>
---
 db.go      |  2 +-
 db_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/db.go b/db.go
index 729e619f5..d30ac6c3d 100644
--- a/db.go
+++ b/db.go
@@ -546,6 +546,7 @@ func (db *DB) close() (err error) {
 	db.opt.Infof("Lifetime L0 stalled for: %s\n", time.Duration(db.lc.l0stallsMs.Load()))
 
 	db.blockWrites.Store(1)
+	db.isClosed.Store(1)
 
 	if !db.opt.InMemory {
 		// Stop value GC first.
@@ -631,7 +632,6 @@ func (db *DB) close() (err error) {
 	db.blockCache.Close()
 	db.indexCache.Close()
 
-	db.isClosed.Store(1)
 	db.threshold.close()
 
 	if db.opt.InMemory {
diff --git a/db_test.go b/db_test.go
index f1933b07b..01b28332a 100644
--- a/db_test.go
+++ b/db_test.go
@@ -2635,3 +2635,39 @@ func TestCompactL0OnClose(t *testing.T) {
 		}
 	})
 }
+
+func TestConcurrent(t *testing.T) {
+	dir, err := os.MkdirTemp("", "badger-concurrent-test")
+	require.NoError(t, err)
+	db, err := Open(DefaultOptions(dir))
+	require.NoError(t, err)
+
+	key := []byte("key")
+	err = db.Update(func(txn *Txn) error {
+		return txn.Set(key, []byte("value"))
+	})
+	require.NoError(t, err)
+
+	var wg sync.WaitGroup
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for {
+				err := db.View(func(txn *Txn) error {
+					_, err := txn.Get(key)
+					return err
+				})
+				if err != nil {
+					require.Contains(t, err.Error(), "DB Closed")
+					break
+				}
+			}
+		}()
+	}
+
+	time.Sleep(time.Second)
+	err = db.Close()
+	require.NoError(t, err)
+	wg.Wait()
+}