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

make stmt cache size configurable

This commit is contained in:
jianzhiyao 2020-07-28 12:57:19 +08:00
parent 93eb7c6b83
commit 756df9385f
4 changed files with 144 additions and 50 deletions

View File

@ -18,4 +18,5 @@ const (
MaxIdleConnsKey = "MaxIdleConns" MaxIdleConnsKey = "MaxIdleConns"
MaxOpenConnsKey = "MaxOpenConns" MaxOpenConnsKey = "MaxOpenConns"
ConnMaxLifetimeKey = "ConnMaxLifetime" ConnMaxLifetimeKey = "ConnMaxLifetime"
MaxStmtCacheSize = "MaxStmtCacheSize"
) )

View File

@ -111,6 +111,7 @@ type DB struct {
*sync.RWMutex *sync.RWMutex
DB *sql.DB DB *sql.DB
stmtDecorators *lru.Cache stmtDecorators *lru.Cache
stmtDecoratorsLimit int
} }
var _ dbQuerier = new(DB) var _ dbQuerier = new(DB)
@ -165,16 +166,14 @@ 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) {
sd, err := d.getStmtDecorator(query) return d.ExecContext(context.Background(), query, args...)
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) { func (d *DB) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
if d.stmtDecorators == nil {
return d.DB.ExecContext(ctx, query, args...)
}
sd, err := d.getStmtDecorator(query) sd, err := d.getStmtDecorator(query)
if err != nil { if err != nil {
return nil, err return nil, err
@ -185,16 +184,14 @@ func (d *DB) ExecContext(ctx context.Context, query string, args ...interface{})
} }
func (d *DB) Query(query string, args ...interface{}) (*sql.Rows, error) { func (d *DB) Query(query string, args ...interface{}) (*sql.Rows, error) {
sd, err := d.getStmtDecorator(query) return d.QueryContext(context.Background(), query, args...)
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) { func (d *DB) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
if d.stmtDecorators == nil {
return d.DB.QueryContext(ctx, query, args...)
}
sd, err := d.getStmtDecorator(query) sd, err := d.getStmtDecorator(query)
if err != nil { if err != nil {
return nil, err return nil, err
@ -205,24 +202,21 @@ func (d *DB) QueryContext(ctx context.Context, query string, args ...interface{}
} }
func (d *DB) QueryRow(query string, args ...interface{}) *sql.Row { func (d *DB) QueryRow(query string, args ...interface{}) *sql.Row {
sd, err := d.getStmtDecorator(query) return d.QueryRowContext(context.Background(), query, args...)
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 { func (d *DB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
if d.stmtDecorators == nil {
return d.DB.QueryRowContext(ctx, query, args...)
}
sd, err := d.getStmtDecorator(query) sd, err := d.getStmtDecorator(query)
if err != nil { if err != nil {
panic(err) panic(err)
} }
stmt := sd.getStmt() stmt := sd.getStmt()
defer sd.release() defer sd.release()
return stmt.QueryRowContext(ctx, args) return stmt.QueryRowContext(ctx, args...)
} }
type TxDB struct { type TxDB struct {
@ -345,14 +339,31 @@ func detectTZ(al *alias) {
} }
} }
func addAliasWthDB(aliasName, driverName string, db *sql.DB) (*alias, error) { func addAliasWthDB(aliasName, driverName string, db *sql.DB, params ...common.KV) (*alias, error) {
kvs := common.NewKVs(params...)
var stmtCache *lru.Cache
var stmtCacheSize int
maxStmtCacheSize := kvs.GetValueOr(MaxStmtCacheSize, 0).(int)
if maxStmtCacheSize > 0 {
_stmtCache, errC := newStmtDecoratorLruWithEvict(maxStmtCacheSize)
if errC != nil {
return nil, errC
} else {
stmtCache = _stmtCache
stmtCacheSize = maxStmtCacheSize
}
}
al := new(alias) al := new(alias)
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,
stmtDecorators: newStmtDecoratorLruWithEvict(), stmtDecorators: stmtCache,
stmtDecoratorsLimit: stmtCacheSize,
} }
if dr, ok := drivers[driverName]; ok { if dr, ok := drivers[driverName]; ok {
@ -371,12 +382,22 @@ func addAliasWthDB(aliasName, driverName string, db *sql.DB) (*alias, error) {
return nil, fmt.Errorf("DataBase alias name `%s` already registered, cannot reuse", aliasName) return nil, fmt.Errorf("DataBase alias name `%s` already registered, cannot reuse", aliasName)
} }
detectTZ(al)
kvs.IfContains(MaxIdleConnsKey, func(value interface{}) {
SetMaxIdleConns(al.Name, value.(int))
}).IfContains(MaxOpenConnsKey, func(value interface{}) {
SetMaxOpenConns(al.Name, value.(int))
}).IfContains(ConnMaxLifetimeKey, func(value interface{}) {
SetConnMaxLifetime(al.Name, value.(time.Duration))
})
return al, nil return al, nil
} }
// AddAliasWthDB add a aliasName for the drivename // AddAliasWthDB add a aliasName for the drivename
func AddAliasWthDB(aliasName, driverName string, db *sql.DB) error { func AddAliasWthDB(aliasName, driverName string, db *sql.DB, params ...common.KV) error {
_, err := addAliasWthDB(aliasName, driverName, db) _, err := addAliasWthDB(aliasName, driverName, db, params...)
return err return err
} }
@ -388,7 +409,6 @@ func RegisterDataBase(aliasName, driverName, dataSource string, params ...common
al *alias al *alias
) )
kvs := common.NewKVs(params...)
db, err = sql.Open(driverName, dataSource) db, err = sql.Open(driverName, dataSource)
if err != nil { if err != nil {
@ -396,23 +416,13 @@ func RegisterDataBase(aliasName, driverName, dataSource string, params ...common
goto end goto end
} }
al, err = addAliasWthDB(aliasName, driverName, db) al, err = addAliasWthDB(aliasName, driverName, db, params...)
if err != nil { if err != nil {
goto end goto end
} }
al.DataSource = dataSource al.DataSource = dataSource
detectTZ(al)
kvs.IfContains(MaxIdleConnsKey, func(value interface{}) {
SetMaxIdleConns(al.Name, value.(int))
}).IfContains(MaxOpenConnsKey, func(value interface{}) {
SetMaxOpenConns(al.Name, value.(int))
}).IfContains(ConnMaxLifetimeKey, func(value interface{}) {
SetConnMaxLifetime(al.Name, value.(time.Duration))
})
end: end:
if err != nil { if err != nil {
if db != nil { if db != nil {
@ -517,9 +527,12 @@ func newStmtDecorator(sqlStmt *sql.Stmt) *stmtDecorator {
} }
} }
func newStmtDecoratorLruWithEvict() *lru.Cache { func newStmtDecoratorLruWithEvict(cacheSize int) (*lru.Cache, error) {
cache, _ := lru.NewWithEvict(1000, func(key interface{}, value interface{}) { cache, err := lru.NewWithEvict(cacheSize, func(key interface{}, value interface{}) {
value.(*stmtDecorator).destroy() value.(*stmtDecorator).destroy()
}) })
return cache if err != nil {
return nil, err
}
return cache, nil
} }

View File

@ -42,3 +42,56 @@ func TestRegisterDataBase(t *testing.T) {
assert.Equal(t, al.MaxOpenConns, 300) assert.Equal(t, al.MaxOpenConns, 300)
assert.Equal(t, al.ConnMaxLifetime, time.Minute) assert.Equal(t, al.ConnMaxLifetime, time.Minute)
} }
func TestRegisterDataBase_MaxStmtCacheSizeNegative1(t *testing.T) {
aliasName := "TestRegisterDataBase_MaxStmtCacheSizeNegative1"
err := RegisterDataBase(aliasName, DBARGS.Driver, DBARGS.Source, common.KV{
Key: MaxStmtCacheSize,
Value: -1,
})
assert.Nil(t, err)
al := getDbAlias(aliasName)
assert.NotNil(t, al)
assert.Equal(t, al.DB.stmtDecoratorsLimit, 0)
}
func TestRegisterDataBase_MaxStmtCacheSize0(t *testing.T) {
aliasName := "TestRegisterDataBase_MaxStmtCacheSize0"
err := RegisterDataBase(aliasName, DBARGS.Driver, DBARGS.Source, common.KV{
Key: MaxStmtCacheSize,
Value: 0,
})
assert.Nil(t, err)
al := getDbAlias(aliasName)
assert.NotNil(t, al)
assert.Equal(t, al.DB.stmtDecoratorsLimit, 0)
}
func TestRegisterDataBase_MaxStmtCacheSize1(t *testing.T) {
aliasName := "TestRegisterDataBase_MaxStmtCacheSize1"
err := RegisterDataBase(aliasName, DBARGS.Driver, DBARGS.Source, common.KV{
Key: MaxStmtCacheSize,
Value: 1,
})
assert.Nil(t, err)
al := getDbAlias(aliasName)
assert.NotNil(t, al)
assert.Equal(t, al.DB.stmtDecoratorsLimit, 1)
}
func TestRegisterDataBase_MaxStmtCacheSize841(t *testing.T) {
aliasName := "TestRegisterDataBase_MaxStmtCacheSize841"
err := RegisterDataBase(aliasName, DBARGS.Driver, DBARGS.Source, common.KV{
Key: MaxStmtCacheSize,
Value: 841,
})
assert.Nil(t, err)
al := getDbAlias(aliasName)
assert.NotNil(t, al)
assert.Equal(t, al.DB.stmtDecoratorsLimit, 841)
}

View File

@ -58,6 +58,8 @@ import (
"database/sql" "database/sql"
"errors" "errors"
"fmt" "fmt"
"github.com/astaxie/beego/pkg/common"
lru "github.com/hashicorp/golang-lru"
"os" "os"
"reflect" "reflect"
"sync" "sync"
@ -609,7 +611,7 @@ func NewOrm() Ormer {
} }
// NewOrmWithDB create a new ormer object with specify *sql.DB for query // NewOrmWithDB create a new ormer object with specify *sql.DB for query
func NewOrmWithDB(driverName, aliasName string, db *sql.DB) (Ormer, error) { func NewOrmWithDB(driverName, aliasName string, db *sql.DB, params ...common.KV) (Ormer, error) {
var al *alias var al *alias
if dr, ok := drivers[driverName]; ok { if dr, ok := drivers[driverName]; ok {
@ -620,16 +622,41 @@ func NewOrmWithDB(driverName, aliasName string, db *sql.DB) (Ormer, error) {
return nil, fmt.Errorf("driver name `%s` have not registered", driverName) return nil, fmt.Errorf("driver name `%s` have not registered", driverName)
} }
kvs := common.NewKVs(params...)
var stmtCache *lru.Cache
var stmtCacheSize int
maxStmtCacheSize := kvs.GetValueOr(MaxStmtCacheSize, 0).(int)
if maxStmtCacheSize > 0 {
_stmtCache, errC := newStmtDecoratorLruWithEvict(maxStmtCacheSize)
if errC != nil {
return nil, errC
} else {
stmtCache = _stmtCache
stmtCacheSize = maxStmtCacheSize
}
}
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,
stmtDecorators: newStmtDecoratorLruWithEvict(), stmtDecorators: stmtCache,
stmtDecoratorsLimit: stmtCacheSize,
} }
detectTZ(al) detectTZ(al)
kvs.IfContains(MaxIdleConnsKey, func(value interface{}) {
SetMaxIdleConns(al.Name, value.(int))
}).IfContains(MaxOpenConnsKey, func(value interface{}) {
SetMaxOpenConns(al.Name, value.(int))
}).IfContains(ConnMaxLifetimeKey, func(value interface{}) {
SetConnMaxLifetime(al.Name, value.(time.Duration))
})
o := new(orm) o := new(orm)
o.alias = al o.alias = al