Repository: dgraph-io/badger
Pull request #1480 (2020-08-26 10:05:38)

release/v2.2007 - Add IsClosed method (#1475)
Add an IsClosed method which denotes if badger instance is closed or not.

(cherry picked from commit 1e21a9449ab01a9667a9c1fa84d9a0b70bef4fd8)

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/dgraph-io/badger/1480)
<!-- Reviewable:end -->



Patch:
diff --git a/db.go b/db.go
index 6818661..07945c8 100644
--- a/db.go
+++ b/db.go
@@ -84,6 +84,7 @@ type DB struct {
 	logRotates int32
 
 	blockWrites int32
+	isClosed    uint32
 
 	orc *oracle
 
@@ -474,6 +475,12 @@ func (db *DB) Close() error {
 	return err
 }
 
+// IsClosed denotes if the badger DB is closed or not. A DB instance should not
+// be used after closing it.
+func (db *DB) IsClosed() bool {
+	return atomic.LoadUint32(&db.isClosed) == 1
+}
+
 func (db *DB) close() (err error) {
 	db.opt.Debugf("Closing database")
 
@@ -555,6 +562,8 @@ func (db *DB) close() (err error) {
 	db.blockCache.Close()
 	db.indexCache.Close()
 
+	atomic.StoreUint32(&db.isClosed, 1)
+
 	if db.opt.InMemory {
 		return
 	}
@@ -646,6 +655,9 @@ func (db *DB) getMemTables() ([]*skl.Skiplist, func()) {
 // been moved, then for the corresponding movekey, we'll look through all the levels of the tree
 // to ensure that we pick the highest version of the movekey present.
 func (db *DB) get(key []byte) (y.ValueStruct, error) {
+	if db.IsClosed() {
+		return y.ValueStruct{}, ErrDBClosed
+	}
 	tables, decr := db.getMemTables() // Lock should be released.
 	defer decr()
 
diff --git a/errors.go b/errors.go
index 61dec9b..65095f1 100644
--- a/errors.go
+++ b/errors.go
@@ -117,8 +117,12 @@ var (
 	// ErrInvalidDataKeyID is returned if the datakey id is invalid.
 	ErrInvalidDataKeyID = errors.New("Invalid datakey id")
 
+	// ErrInvalidEncryptionKey is returned if length of encryption keys is invalid.
 	ErrInvalidEncryptionKey = errors.New("Encryption key's length should be" +
 		"either 16, 24, or 32 bytes")
-
+	// ErrGCInMemoryMode is returned when db.RunValueLogGC is called in in-memory mode.
 	ErrGCInMemoryMode = errors.New("Cannot run value log GC when DB is opened in InMemory mode")
+
+	// ErrDBClosed is returned when a get operation is performed after closing the DB.
+	ErrDBClosed = errors.New("DB Closed")
 )
diff --git a/iterator.go b/iterator.go
index 887b049..11d0c27 100644
--- a/iterator.go
+++ b/iterator.go
@@ -459,6 +459,9 @@ func (txn *Txn) NewIterator(opt IteratorOptions) *Iterator {
 	if txn.discarded {
 		panic("Transaction has already been discarded")
 	}
+	if txn.db.IsClosed() {
+		panic(ErrDBClosed.Error())
+	}
 
 	// Keep track of the number of active iterators.
 	atomic.AddInt32(&txn.numIterators, 1)
diff --git a/levels.go b/levels.go
index aaef7e3..f99ba0a 100644
--- a/levels.go
+++ b/levels.go
@@ -1073,6 +1073,9 @@ func (s *levelsController) close() error {
 // get returns the found value if any. If not found, we return nil.
 func (s *levelsController) get(key []byte, maxVs *y.ValueStruct, startLevel int) (
 	y.ValueStruct, error) {
+	if s.kv.IsClosed() {
+		return y.ValueStruct{}, ErrDBClosed
+	}
 	// It's important that we iterate the levels from 0 on upward. The reason is, if we iterated
 	// in opposite order, or in parallel (naively calling all the h.RLock() in some order) we could
 	// read level L's tables post-compaction and level L+1's tables pre-compaction. (If we do
diff --git a/txn.go b/txn.go
index 3b16cb9..8fc5381 100644
--- a/txn.go
+++ b/txn.go
@@ -788,6 +788,9 @@ func (db *DB) newTransaction(update, isManaged bool) *Txn {
 // returned by the function is relayed by the View method.
 // If View is used with managed transactions, it would assume a read timestamp of MaxUint64.
 func (db *DB) View(fn func(txn *Txn) error) error {
+	if db.IsClosed() {
+		return ErrDBClosed
+	}
 	var txn *Txn
 	if db.opt.managedTxns {
 		txn = db.NewTransactionAt(math.MaxUint64, false)
@@ -803,6 +806,9 @@ func (db *DB) View(fn func(txn *Txn) error) error {
 // for the user. Error returned by the function is relayed by the Update method.
 // Update cannot be used with managed transactions.
 func (db *DB) Update(fn func(txn *Txn) error) error {
+	if db.IsClosed() {
+		return ErrDBClosed
+	}
 	if db.opt.managedTxns {
 		panic("Update can only be used with managedDB=false.")
 	}


Test patch:
diff --git a/db2_test.go b/db2_test.go
index 4c29d63..645971b 100644
--- a/db2_test.go
+++ b/db2_test.go
@@ -831,3 +831,33 @@ func TestDropAllDropPrefix(t *testing.T) {
 		wg.Wait()
 	})
 }
+
+func TestIsClosed(t *testing.T) {
+	test := func(inMemory bool) {
+		opt := DefaultOptions("")
+		if inMemory {
+			opt.InMemory = true
+		} else {
+			dir, err := ioutil.TempDir("", "badger-test")
+			require.NoError(t, err)
+			defer removeDir(dir)
+
+			opt.Dir = dir
+			opt.ValueDir = dir
+		}
+
+		db, err := Open(opt)
+		require.NoError(t, err)
+		require.False(t, db.IsClosed())
+		require.NoError(t, db.Close())
+		require.True(t, db.IsClosed())
+	}
+
+	t.Run("normal", func(t *testing.T) {
+		test(false)
+	})
+	t.Run("in-memory", func(t *testing.T) {
+		test(true)
+	})
+
+}
diff --git a/managed_db_test.go b/managed_db_test.go
index cd5bd15..58165d7 100644
--- a/managed_db_test.go
+++ b/managed_db_test.go
@@ -306,7 +306,7 @@ func TestWriteAfterClose(t *testing.T) {
 	err = db.Update(func(txn *Txn) error {
 		return txn.SetEntry(NewEntry([]byte("a"), []byte("b")))
 	})
-	require.Equal(t, ErrBlockedWrites, err)
+	require.Equal(t, ErrDBClosed, err)
 }
 
 func TestDropAllRace(t *testing.T) {