fix: prepare deadlock

This commit is contained in:
a631807682 2022-08-02 14:23:56 +08:00
parent 6e03b97e26
commit b73096d61d
No known key found for this signature in database
GPG Key ID: 137D1D75522168AB
2 changed files with 34 additions and 14 deletions

View File

@ -50,23 +50,18 @@ func (db *PreparedStmtDB) prepare(ctx context.Context, conn ConnPool, isTransact
} }
db.Mux.RUnlock() db.Mux.RUnlock()
db.Mux.Lock()
defer db.Mux.Unlock()
// double check
if stmt, ok := db.Stmts[query]; ok && (!stmt.Transaction || isTransaction) {
return stmt, nil
} else if ok {
go stmt.Close()
}
stmt, err := conn.PrepareContext(ctx, query) stmt, err := conn.PrepareContext(ctx, query)
if err == nil { if err != nil {
db.Stmts[query] = Stmt{Stmt: stmt, Transaction: isTransaction} return Stmt{}, err
db.PreparedSQL = append(db.PreparedSQL, query)
} }
return db.Stmts[query], err cacheStmt := Stmt{Stmt: stmt, Transaction: isTransaction}
db.Mux.Lock()
db.Stmts[query] = cacheStmt
db.PreparedSQL = append(db.PreparedSQL, query)
db.Mux.Unlock()
return cacheStmt, err
} }
func (db *PreparedStmtDB) BeginTx(ctx context.Context, opt *sql.TxOptions) (ConnPool, error) { func (db *PreparedStmtDB) BeginTx(ctx context.Context, opt *sql.TxOptions) (ConnPool, error) {

View File

@ -2,6 +2,7 @@ package tests_test
import ( import (
"context" "context"
"sync"
"testing" "testing"
"time" "time"
@ -88,3 +89,27 @@ func TestPreparedStmtFromTransaction(t *testing.T) {
} }
tx2.Commit() tx2.Commit()
} }
func TestPreparedStmtDeadlock(t *testing.T) {
tx, err := OpenTestConnection()
AssertEqual(t, err, nil)
sqlDB, _ := tx.DB()
sqlDB.SetMaxOpenConns(1)
tx = tx.Session(&gorm.Session{PrepareStmt: true})
wg := sync.WaitGroup{}
for i := 0; i < 2; i++ {
wg.Add(1)
go func(j int) {
user := User{Name: "jinzhu"}
tx.Create(&user)
var result User
tx.First(&result)
wg.Done()
}(i)
}
wg.Wait()
}