1
0
mirror of https://github.com/astaxie/beego.git synced 2024-11-22 08:30:54 +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" "context"
"database/sql" "database/sql"
"fmt" "fmt"
lru "github.com/hashicorp/golang-lru"
"reflect" "reflect"
"sync" "sync"
"time" "time"
@ -106,8 +107,8 @@ func (ac *_dbCache) getDefault() (al *alias) {
type DB struct { type DB struct {
*sync.RWMutex *sync.RWMutex
DB *sql.DB DB *sql.DB
stmts map[string]*sql.Stmt stmtDecorators *lru.Cache
} }
func (d *DB) Begin() (*sql.Tx, error) { 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) 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() d.RLock()
c, ok := d.stmts[query] c, ok := d.stmtDecorators.Get(query)
d.RUnlock() d.RUnlock()
if ok { if ok {
return c, nil c.(*stmtDecorator).acquire()
return c.(*stmtDecorator), nil
} }
d.Lock() d.Lock()
c, ok = d.stmts[query] c, ok = d.stmtDecorators.Get(query)
if ok { if ok {
d.Unlock() d.Unlock()
return c, nil c.(*stmtDecorator).acquire()
return c.(*stmtDecorator), nil
} }
stmt, err := d.Prepare(query) stmt, err := d.Prepare(query)
@ -138,9 +142,12 @@ func (d *DB) getStmt(query string) (*sql.Stmt, error) {
d.Unlock() d.Unlock()
return nil, err return nil, err
} }
d.stmts[query] = stmt sd := newStmtDecorator(stmt)
d.stmtDecorators.Add(query, sd)
d.Unlock() d.Unlock()
return stmt, nil
sd.acquire()
return sd, nil
} }
func (d *DB) Prepare(query string) (*sql.Stmt, error) { 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) { func (d *DB) Exec(query string, args ...interface{}) (sql.Result, error) {
stmt, err := d.getStmt(query) sd, err := d.getStmtDecorator(query)
if err != nil { if err != nil {
return nil, err return nil, err
} }
stmt := sd.getStmt()
defer sd.release()
return stmt.Exec(args...) return stmt.Exec(args...)
} }
func (d *DB) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) { 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 { if err != nil {
return nil, err return nil, err
} }
stmt := sd.getStmt()
defer sd.release()
return stmt.ExecContext(ctx, args...) return stmt.ExecContext(ctx, args...)
} }
func (d *DB) Query(query string, args ...interface{}) (*sql.Rows, error) { func (d *DB) Query(query string, args ...interface{}) (*sql.Rows, error) {
stmt, err := d.getStmt(query) sd, err := d.getStmtDecorator(query)
if err != nil { if err != nil {
return nil, err return nil, err
} }
stmt := sd.getStmt()
defer sd.release()
return stmt.Query(args...) return stmt.Query(args...)
} }
func (d *DB) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) { 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 { if err != nil {
return nil, err return nil, err
} }
stmt := sd.getStmt()
defer sd.release()
return stmt.QueryContext(ctx, args...) return stmt.QueryContext(ctx, args...)
} }
func (d *DB) QueryRow(query string, args ...interface{}) *sql.Row { func (d *DB) QueryRow(query string, args ...interface{}) *sql.Row {
stmt, err := d.getStmt(query) sd, err := d.getStmtDecorator(query)
if err != nil { if err != nil {
panic(err) panic(err)
} }
stmt := sd.getStmt()
defer sd.release()
return stmt.QueryRow(args...) return stmt.QueryRow(args...)
} }
func (d *DB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row { func (d *DB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
sd, err := d.getStmtDecorator(query)
stmt, err := d.getStmt(query)
if err != nil { if err != nil {
panic(err) panic(err)
} }
stmt := sd.getStmt()
defer sd.release()
return stmt.QueryRowContext(ctx, args) return stmt.QueryRowContext(ctx, args)
} }
@ -275,9 +293,9 @@ func addAliasWthDB(aliasName, driverName string, db *sql.DB) (*alias, error) {
al.Name = aliasName al.Name = aliasName
al.DriverName = driverName al.DriverName = driverName
al.DB = &DB{ al.DB = &DB{
RWMutex: new(sync.RWMutex), RWMutex: new(sync.RWMutex),
DB: db, DB: db,
stmts: make(map[string]*sql.Stmt), stmtDecorators: newStmtDecoratorLruWithEvict(),
} }
if dr, ok := drivers[driverName]; ok { 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) 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.Name = aliasName
al.DriverName = driverName al.DriverName = driverName
al.DB = &DB{ al.DB = &DB{
RWMutex: new(sync.RWMutex), RWMutex: new(sync.RWMutex),
DB: db, DB: db,
stmts: make(map[string]*sql.Stmt), stmtDecorators: newStmtDecoratorLruWithEvict(),
} }
detectTZ(al) detectTZ(al)