1
0
mirror of https://github.com/astaxie/beego.git synced 2024-11-22 02:30:53 +00:00

add stmt garbage recycle

This commit is contained in:
jianzhiyao 2020-06-22 23:23:52 +08:00
parent decc75e025
commit a5e8344a0a
2 changed files with 76 additions and 22 deletions

View File

@ -18,6 +18,7 @@ import (
"context"
"database/sql"
"fmt"
lru "github.com/hashicorp/golang-lru"
"reflect"
"sync"
"time"
@ -106,8 +107,8 @@ func (ac *_dbCache) getDefault() (al *alias) {
type DB struct {
*sync.RWMutex
DB *sql.DB
stmts map[string]*sql.Stmt
DB *sql.DB
stmtDecorators *lru.Cache
}
func (d *DB) Begin() (*sql.Tx, error) {
@ -118,19 +119,22 @@ func (d *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error)
return d.DB.BeginTx(ctx, opts)
}
func (d *DB) getStmt(query string) (*sql.Stmt, error) {
//su must call release to release *sql.Stmt after using
func (d *DB) getStmtDecorator(query string) (*stmtDecorator, error) {
d.RLock()
c, ok := d.stmts[query]
c, ok := d.stmtDecorators.Get(query)
d.RUnlock()
if ok {
return c, nil
c.(*stmtDecorator).acquire()
return c.(*stmtDecorator), nil
}
d.Lock()
c, ok = d.stmts[query]
c, ok = d.stmtDecorators.Get(query)
if ok {
d.Unlock()
return c, nil
c.(*stmtDecorator).acquire()
return c.(*stmtDecorator), nil
}
stmt, err := d.Prepare(query)
@ -138,9 +142,12 @@ func (d *DB) getStmt(query string) (*sql.Stmt, error) {
d.Unlock()
return nil, err
}
d.stmts[query] = stmt
sd := newStmtDecorator(stmt)
d.stmtDecorators.Add(query, sd)
d.Unlock()
return stmt, nil
sd.acquire()
return sd, nil
}
func (d *DB) Prepare(query string) (*sql.Stmt, error) {
@ -152,52 +159,63 @@ func (d *DB) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error
}
func (d *DB) Exec(query string, args ...interface{}) (sql.Result, error) {
stmt, err := d.getStmt(query)
sd, err := d.getStmtDecorator(query)
if err != nil {
return nil, err
}
stmt := sd.getStmt()
defer sd.release()
return stmt.Exec(args...)
}
func (d *DB) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
stmt, err := d.getStmt(query)
sd, err := d.getStmtDecorator(query)
if err != nil {
return nil, err
}
stmt := sd.getStmt()
defer sd.release()
return stmt.ExecContext(ctx, args...)
}
func (d *DB) Query(query string, args ...interface{}) (*sql.Rows, error) {
stmt, err := d.getStmt(query)
sd, err := d.getStmtDecorator(query)
if err != nil {
return nil, err
}
stmt := sd.getStmt()
defer sd.release()
return stmt.Query(args...)
}
func (d *DB) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
stmt, err := d.getStmt(query)
sd, err := d.getStmtDecorator(query)
if err != nil {
return nil, err
}
stmt := sd.getStmt()
defer sd.release()
return stmt.QueryContext(ctx, args...)
}
func (d *DB) QueryRow(query string, args ...interface{}) *sql.Row {
stmt, err := d.getStmt(query)
sd, err := d.getStmtDecorator(query)
if err != nil {
panic(err)
}
stmt := sd.getStmt()
defer sd.release()
return stmt.QueryRow(args...)
}
func (d *DB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
stmt, err := d.getStmt(query)
sd, err := d.getStmtDecorator(query)
if err != nil {
panic(err)
}
stmt := sd.getStmt()
defer sd.release()
return stmt.QueryRowContext(ctx, args)
}
@ -275,9 +293,9 @@ func addAliasWthDB(aliasName, driverName string, db *sql.DB) (*alias, error) {
al.Name = aliasName
al.DriverName = driverName
al.DB = &DB{
RWMutex: new(sync.RWMutex),
DB: db,
stmts: make(map[string]*sql.Stmt),
RWMutex: new(sync.RWMutex),
DB: db,
stmtDecorators: newStmtDecoratorLruWithEvict(),
}
if dr, ok := drivers[driverName]; ok {
@ -402,3 +420,39 @@ func GetDB(aliasNames ...string) (*sql.DB, error) {
}
return nil, fmt.Errorf("DataBase of alias name `%s` not found", name)
}
type stmtDecorator struct {
wg sync.WaitGroup
stmt *sql.Stmt
}
func (s *stmtDecorator) getStmt() *sql.Stmt {
return s.stmt
}
func (s *stmtDecorator) acquire() {
s.wg.Add(1)
}
func (s *stmtDecorator) release() {
s.wg.Done()
}
//garbage recycle for stmt
func (s *stmtDecorator) destroy() {
s.wg.Wait()
_ = s.stmt.Close()
}
func newStmtDecorator(sqlStmt *sql.Stmt) *stmtDecorator {
return &stmtDecorator{
stmt: sqlStmt,
}
}
func newStmtDecoratorLruWithEvict() *lru.Cache {
cache, _ := lru.NewWithEvict(1000, func(key interface{}, value interface{}) {
value.(*stmtDecorator).destroy()
})
return cache
}

View File

@ -559,9 +559,9 @@ func NewOrmWithDB(driverName, aliasName string, db *sql.DB) (Ormer, error) {
al.Name = aliasName
al.DriverName = driverName
al.DB = &DB{
RWMutex: new(sync.RWMutex),
DB: db,
stmts: make(map[string]*sql.Stmt),
RWMutex: new(sync.RWMutex),
DB: db,
stmtDecorators: newStmtDecoratorLruWithEvict(),
}
detectTZ(al)