fix txnstore data race (#4579)

fix data race in txnstore.getOrSetDB. add dbsMu

Approved by: @XuPeng-SH

pkg/vm/engine/tae/txn/txnimpl/store.go

@@ -15,6 +15,7 @@
 package txnimpl
 
 import (
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -34,6 +35,7 @@ import (
 type txnStore struct {
 	txnbase.NoopTxnStore
 	dbs         map[uint64]*txnDB
+	mu          sync.RWMutex
 	driver      wal.Driver
 	nodesMgr    base.INodeManager
 	txn         txnif.AsyncTxn
@@ -284,16 +286,25 @@ func (store *txnStore) CreateNonAppendableSegment(dbId, tid uint64) (seg handle.
 }
 
 func (store *txnStore) getOrSetDB(id uint64) (db *txnDB, err error) {
+	store.mu.RLock()
 	db = store.dbs[id]
-	if db == nil {
-		var entry *catalog.DBEntry
-		if entry, err = store.catalog.GetDatabaseByID(id); err != nil {
-			return
-		}
-		db = newTxnDB(store, entry)
-		db.idx = len(store.dbs)
-		store.dbs[id] = db
+	store.mu.RUnlock()
+	if db != nil {
+		return
+	}
+	var entry *catalog.DBEntry
+	if entry, err = store.catalog.GetDatabaseByID(id); err != nil {
+		return
+	}
+	store.mu.Lock()
+	defer store.mu.Unlock()
+	db = store.dbs[id]
+	if db != nil {
+		return
 	}
+	db = newTxnDB(store, entry)
+	db.idx = len(store.dbs)
+	store.dbs[id] = db
 	return
 }
 

pkg/vm/engine/tae/txn/txnimpl/txndb.go

@@ -15,6 +15,7 @@
 package txnimpl
 
 import (
+	"sync"
 	"time"
 
 	"github.com/matrixorigin/matrixone/pkg/logutil"
@@ -29,6 +30,7 @@ import (
 type txnDB struct {
 	store       *txnStore
 	tables      map[uint64]*txnTable
+	mu          sync.RWMutex
 	entry       *catalog.DBEntry
 	createEntry txnif.TxnEntry
 	dropEntry   txnif.TxnEntry
@@ -234,19 +236,28 @@ func (db *txnDB) CreateNonAppendableSegment(tid uint64) (seg handle.Segment, err
 }
 
 func (db *txnDB) getOrSetTable(id uint64) (table *txnTable, err error) {
+	db.mu.RLock()
 	table = db.tables[id]
-	if table == nil {
-		var entry *catalog.TableEntry
-		if entry, err = db.entry.GetTableEntryByID(id); err != nil {
-			return
-		}
-		if db.store.warChecker == nil {
-			db.store.warChecker = newWarChecker(db.store.txn, db.store.catalog)
-		}
-		table = newTxnTable(db.store, entry)
-		table.idx = len(db.tables)
-		db.tables[id] = table
+	db.mu.RUnlock()
+	if table != nil {
+		return
+	}
+	var entry *catalog.TableEntry
+	if entry, err = db.entry.GetTableEntryByID(id); err != nil {
+		return
+	}
+	db.mu.Lock()
+	defer db.mu.Unlock()
+	table = db.tables[id]
+	if table != nil {
+		return
+	}
+	if db.store.warChecker == nil {
+		db.store.warChecker = newWarChecker(db.store.txn, db.store.catalog)
 	}
+	table = newTxnTable(db.store, entry)
+	table.idx = len(db.tables)
+	db.tables[id] = table
 	return
 }