Merge pull request #4028 from jianzhiyao/develop

acquire() in Lock
This commit is contained in:
Ming Deng 2020-06-23 23:58:42 +08:00 committed by GitHub
commit 6052524a28
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 20 additions and 11 deletions

View File

@ -123,14 +123,17 @@ func (d *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error)
func (d *DB) getStmtDecorator(query string) (*stmtDecorator, error) { func (d *DB) getStmtDecorator(query string) (*stmtDecorator, error) {
d.RLock() d.RLock()
c, ok := d.stmtDecorators.Get(query) c, ok := d.stmtDecorators.Get(query)
d.RUnlock()
if ok { if ok {
c.(*stmtDecorator).acquire()
d.RUnlock()
return c.(*stmtDecorator), nil return c.(*stmtDecorator), nil
} }
d.RUnlock()
d.Lock() d.Lock()
c, ok = d.stmtDecorators.Get(query) c, ok = d.stmtDecorators.Get(query)
if ok { if ok {
c.(*stmtDecorator).acquire()
d.Unlock() d.Unlock()
return c.(*stmtDecorator), nil return c.(*stmtDecorator), nil
} }
@ -141,6 +144,7 @@ func (d *DB) getStmtDecorator(query string) (*stmtDecorator, error) {
return nil, err return nil, err
} }
sd := newStmtDecorator(stmt) sd := newStmtDecorator(stmt)
sd.acquire()
d.stmtDecorators.Add(query, sd) d.stmtDecorators.Add(query, sd)
d.Unlock() d.Unlock()
@ -160,7 +164,7 @@ func (d *DB) Exec(query string, args ...interface{}) (sql.Result, error) {
if err != nil { if err != nil {
return nil, err return nil, err
} }
stmt := sd.acquire() stmt := sd.getStmt()
defer sd.release() defer sd.release()
return stmt.Exec(args...) return stmt.Exec(args...)
} }
@ -170,7 +174,7 @@ func (d *DB) ExecContext(ctx context.Context, query string, args ...interface{})
if err != nil { if err != nil {
return nil, err return nil, err
} }
stmt := sd.acquire() stmt := sd.getStmt()
defer sd.release() defer sd.release()
return stmt.ExecContext(ctx, args...) return stmt.ExecContext(ctx, args...)
} }
@ -180,7 +184,7 @@ func (d *DB) Query(query string, args ...interface{}) (*sql.Rows, error) {
if err != nil { if err != nil {
return nil, err return nil, err
} }
stmt := sd.acquire() stmt := sd.getStmt()
defer sd.release() defer sd.release()
return stmt.Query(args...) return stmt.Query(args...)
} }
@ -190,7 +194,7 @@ func (d *DB) QueryContext(ctx context.Context, query string, args ...interface{}
if err != nil { if err != nil {
return nil, err return nil, err
} }
stmt := sd.acquire() stmt := sd.getStmt()
defer sd.release() defer sd.release()
return stmt.QueryContext(ctx, args...) return stmt.QueryContext(ctx, args...)
} }
@ -200,7 +204,7 @@ func (d *DB) QueryRow(query string, args ...interface{}) *sql.Row {
if err != nil { if err != nil {
panic(err) panic(err)
} }
stmt := sd.acquire() stmt := sd.getStmt()
defer sd.release() defer sd.release()
return stmt.QueryRow(args...) return stmt.QueryRow(args...)
@ -211,7 +215,7 @@ func (d *DB) QueryRowContext(ctx context.Context, query string, args ...interfac
if err != nil { if err != nil {
panic(err) panic(err)
} }
stmt := sd.acquire() stmt := sd.getStmt()
defer sd.release() defer sd.release()
return stmt.QueryRowContext(ctx, args) return stmt.QueryRowContext(ctx, args)
} }
@ -425,10 +429,13 @@ type stmtDecorator struct {
stmt *sql.Stmt stmt *sql.Stmt
} }
func (s *stmtDecorator) acquire() *sql.Stmt{ func (s *stmtDecorator) getStmt() *sql.Stmt {
return s.stmt
}
func (s *stmtDecorator) acquire() {
s.wg.Add(1) s.wg.Add(1)
s.lastUse = time.Now().Unix() s.lastUse = time.Now().Unix()
return s.stmt
} }
func (s *stmtDecorator) release() { func (s *stmtDecorator) release() {
@ -437,8 +444,10 @@ func (s *stmtDecorator) release() {
//garbage recycle for stmt //garbage recycle for stmt
func (s *stmtDecorator) destroy() { func (s *stmtDecorator) destroy() {
s.wg.Wait() go func() {
_ = s.stmt.Close() s.wg.Wait()
_ = s.stmt.Close()
}()
} }
func newStmtDecorator(sqlStmt *sql.Stmt) *stmtDecorator { func newStmtDecorator(sqlStmt *sql.Stmt) *stmtDecorator {