Fix: race condition in Database.Get

packages/database/database.go

@@ -10,22 +10,19 @@ var dbMap = make(map[string]*prefixDb)
 var mu sync.Mutex
 
 type prefixDb struct {
-	db       *badger.DB
-	name     string
-	prefix   []byte
-	openLock sync.Mutex
+	db     *badger.DB
+	name   string
+	prefix []byte
 }
 
-func Get(name string) (Database, error) {
-	// avoid locking if it's a clean hit
-	if db, exists := dbMap[name]; exists {
-		return db, nil
-	}
+func getPrefix(name string) []byte {
+	return []byte(name + "_")
+}
 
+func Get(name string) (Database, error) {
 	mu.Lock()
 	defer mu.Unlock()
 
-	// needs to be re-checked after locking
 	if db, exists := dbMap[name]; exists {
 		return db, nil
 	}
@@ -34,7 +31,7 @@ func Get(name string) (Database, error) {
 	db := &prefixDb{
 		db:     badger,
 		name:   name,
-		prefix: []byte(name + "_"),
+		prefix: getPrefix(name),
 	}
 
 	dbMap[name] = db
@@ -43,21 +40,16 @@ func Get(name string) (Database, error) {
 }
 
 func (this *prefixDb) Set(key []byte, value []byte) error {
-	if err := this.db.Update(func(txn *badger.Txn) error { return txn.Set(append(this.prefix, key...), value) }); err != nil {
-		return err
-	}
-
-	return nil
+	err := this.db.Update(func(txn *badger.Txn) error {
+		return txn.Set(append(this.prefix, key...), value)
+	})
+	return err
 }
 
 func (this *prefixDb) Contains(key []byte) (bool, error) {
 	err := this.db.View(func(txn *badger.Txn) error {
 		_, err := txn.Get(append(this.prefix, key...))
-		if err != nil {
-			return err
-		}
-
-		return nil
+		return err
 	})
 
 	if err == badger.ErrKeyNotFound {
@@ -88,8 +80,7 @@ func (this *prefixDb) Get(key []byte) ([]byte, error) {
 
 func (this *prefixDb) Delete(key []byte) error {
 	err := this.db.Update(func(txn *badger.Txn) error {
-		err := txn.Delete(append(this.prefix, key...))
-		return err
+		return txn.Delete(append(this.prefix, key...))
 	})
 	return err
 }