1.4.5.2

EdgeDNS/internal/utils/dbs/batch.go

@@ -0,0 +1,192 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package dbs
+
+import (
+	"database/sql"
+	"github.com/TeaOSLab/EdgeDNS/internal/remotelogs"
+	"time"
+)
+
+type batchItem struct {
+	query string
+	args  []any
+}
+
+type Batch struct {
+	db *DB
+	n  int
+
+	enableStat bool
+
+	onFail func(err error)
+
+	queue      chan *batchItem
+	closeEvent chan bool
+
+	isClosed bool
+}
+
+func NewBatch(db *DB, n int) *Batch {
+	var batch = &Batch{
+		db:         db,
+		n:          n,
+		queue:      make(chan *batchItem, 16),
+		closeEvent: make(chan bool, 1),
+	}
+	db.batches = append(db.batches, batch)
+	return batch
+}
+
+func (this *Batch) EnableStat(b bool) {
+	this.enableStat = b
+}
+
+func (this *Batch) OnFail(callback func(err error)) {
+	this.onFail = callback
+}
+
+func (this *Batch) Add(query string, args ...any) {
+	if this.isClosed {
+		return
+	}
+	this.queue <- &batchItem{
+		query: query,
+		args:  args,
+	}
+}
+
+func (this *Batch) Exec() {
+	var n = this.n
+	if n <= 0 {
+		n = 4
+	}
+
+	var ticker = time.NewTicker(100 * time.Millisecond)
+	var count = 0
+	var lastTx *sql.Tx
+For:
+	for {
+		// closed
+		if this.isClosed {
+			if lastTx != nil {
+				_ = this.commitTx(lastTx)
+				lastTx = nil
+			}
+
+			return
+		}
+
+		select {
+		case item := <-this.queue:
+			if lastTx == nil {
+				lastTx = this.beginTx()
+				if lastTx == nil {
+					continue For
+				}
+			}
+
+			err := this.execItem(lastTx, item)
+			if err != nil {
+				if IsClosedErr(err) {
+					return
+				}
+				this.processErr(item.query, err)
+			}
+
+			count++
+
+			if count == n {
+				count = 0
+				err = this.commitTx(lastTx)
+				lastTx = nil
+				if err != nil {
+					if IsClosedErr(err) {
+						return
+					}
+					this.processErr("commit", err)
+				}
+			}
+		case <-ticker.C:
+			if lastTx == nil || count == 0 {
+				continue For
+			}
+			count = 0
+			err := this.commitTx(lastTx)
+			lastTx = nil
+			if err != nil {
+				if IsClosedErr(err) {
+					return
+				}
+				this.processErr("commit", err)
+			}
+		case <-this.closeEvent:
+			// closed
+
+			if lastTx != nil {
+				_ = this.commitTx(lastTx)
+				lastTx = nil
+			}
+
+			return
+		}
+	}
+}
+
+func (this *Batch) close() {
+	this.isClosed = true
+
+	select {
+	case this.closeEvent <- true:
+	default:
+
+	}
+}
+
+func (this *Batch) beginTx() *sql.Tx {
+	if !this.db.BeginUpdating() {
+		return nil
+	}
+
+	tx, err := this.db.Begin()
+	if err != nil {
+		this.processErr("begin transaction", err)
+		this.db.EndUpdating()
+		return nil
+	}
+	return tx
+}
+
+func (this *Batch) commitTx(tx *sql.Tx) error {
+	// always commit without checking database closing status
+	this.db.EndUpdating()
+	return tx.Commit()
+}
+
+func (this *Batch) execItem(tx *sql.Tx, item *batchItem) error {
+	// check database status
+	if this.db.BeginUpdating() {
+		defer this.db.EndUpdating()
+	} else {
+		return errDBIsClosed
+	}
+
+	if this.enableStat {
+		defer SharedQueryStatManager.AddQuery(item.query).End()
+	}
+
+	_, err := tx.Exec(item.query, item.args...)
+	return err
+}
+
+func (this *Batch) processErr(prefix string, err error) {
+	if err == nil {
+		return
+	}
+
+	if this.onFail != nil {
+		this.onFail(err)
+	} else {
+		remotelogs.Error("SQLITE_BATCH", prefix+": "+err.Error())
+	}
+}

EdgeDNS/internal/utils/dbs/db.go

@@ -0,0 +1,252 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package dbs
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	teaconst "github.com/TeaOSLab/EdgeDNS/internal/const"
+	"github.com/TeaOSLab/EdgeDNS/internal/events"
+	"github.com/TeaOSLab/EdgeDNS/internal/remotelogs"
+	fsutils "github.com/TeaOSLab/EdgeDNS/internal/utils/fs"
+	_ "github.com/mattn/go-sqlite3"
+	"os"
+	"strings"
+	"sync"
+	"time"
+)
+
+var errDBIsClosed = errors.New("the database is closed")
+
+type DB struct {
+	locker *fsutils.Locker
+	rawDB  *sql.DB
+	dsn    string
+
+	statusLocker  sync.Mutex
+	countUpdating int32
+
+	isClosing bool
+
+	enableStat bool
+
+	batches []*Batch
+}
+
+func OpenWriter(dsn string) (*DB, error) {
+	return open(dsn, true)
+}
+
+func OpenReader(dsn string) (*DB, error) {
+	return open(dsn, false)
+}
+
+func open(dsn string, lock bool) (*DB, error) {
+	if teaconst.IsQuiting {
+		return nil, errors.New("can not open database when process is quiting")
+	}
+
+	// decode path
+	var path = dsn
+	var queryIndex = strings.Index(dsn, "?")
+	if queryIndex >= 0 {
+		path = path[:queryIndex]
+	}
+	path = strings.TrimSpace(strings.TrimPrefix(path, "file:"))
+
+	// locker
+	var locker *fsutils.Locker
+	if lock {
+		locker = fsutils.NewLocker(path)
+		err := locker.Lock()
+		if err != nil {
+			remotelogs.Warn("DB", "lock '"+path+"' failed: "+err.Error())
+			locker = nil
+		}
+	}
+
+	// check if closed successfully last time, if not we recover it
+	var walPath = path + "-wal"
+	_, statWalErr := os.Stat(walPath)
+	var shouldRecover = statWalErr == nil
+
+	// open
+	rawDB, err := sql.Open("sqlite3", dsn)
+	if err != nil {
+		return nil, err
+	}
+
+	if shouldRecover {
+		err = rawDB.Close()
+		if err != nil {
+			return nil, err
+		}
+
+		// open again
+		rawDB, err = sql.Open("sqlite3", dsn)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	var db = NewDB(rawDB, dsn)
+	db.locker = locker
+	return db, nil
+}
+
+func NewDB(rawDB *sql.DB, dsn string) *DB {
+	var db = &DB{
+		rawDB: rawDB,
+		dsn:   dsn,
+	}
+
+	events.On(events.EventQuit, func() {
+		_ = db.Close()
+	})
+	events.On(events.EventTerminated, func() {
+		_ = db.Close()
+	})
+
+	return db
+}
+
+func (this *DB) SetMaxOpenConns(n int) {
+	this.rawDB.SetMaxOpenConns(n)
+}
+
+func (this *DB) EnableStat(b bool) {
+	this.enableStat = b
+}
+
+func (this *DB) Begin() (*sql.Tx, error) {
+	// check database status
+	if this.BeginUpdating() {
+		defer this.EndUpdating()
+	} else {
+		return nil, errDBIsClosed
+	}
+
+	return this.rawDB.Begin()
+}
+
+func (this *DB) Prepare(query string) (*Stmt, error) {
+	stmt, err := this.rawDB.Prepare(query)
+	if err != nil {
+		return nil, err
+	}
+
+	var s = NewStmt(this, stmt, query)
+	if this.enableStat {
+		s.EnableStat()
+	}
+	return s, nil
+}
+
+func (this *DB) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
+	// check database status
+	if this.BeginUpdating() {
+		defer this.EndUpdating()
+	} else {
+		return nil, errDBIsClosed
+	}
+
+	if this.enableStat {
+		defer SharedQueryStatManager.AddQuery(query).End()
+	}
+
+	return this.rawDB.ExecContext(ctx, query, args...)
+}
+
+func (this *DB) Exec(query string, args ...any) (sql.Result, error) {
+	// check database status
+	if this.BeginUpdating() {
+		defer this.EndUpdating()
+	} else {
+		return nil, errDBIsClosed
+	}
+
+	if this.enableStat {
+		defer SharedQueryStatManager.AddQuery(query).End()
+	}
+	return this.rawDB.Exec(query, args...)
+}
+
+func (this *DB) Query(query string, args ...any) (*sql.Rows, error) {
+	if this.enableStat {
+		defer SharedQueryStatManager.AddQuery(query).End()
+	}
+	return this.rawDB.Query(query, args...)
+}
+
+func (this *DB) QueryRow(query string, args ...any) *sql.Row {
+	if this.enableStat {
+		defer SharedQueryStatManager.AddQuery(query).End()
+	}
+	return this.rawDB.QueryRow(query, args...)
+}
+
+// Close the database
+func (this *DB) Close() error {
+	// check database status
+	this.statusLocker.Lock()
+	if this.isClosing {
+		this.statusLocker.Unlock()
+		return nil
+	}
+	this.isClosing = true
+	this.statusLocker.Unlock()
+
+	// waiting for updating operations to finish
+	for {
+		this.statusLocker.Lock()
+		var countUpdating = this.countUpdating
+		this.statusLocker.Unlock()
+		if countUpdating <= 0 {
+			break
+		}
+		time.Sleep(1 * time.Millisecond)
+	}
+
+	for _, batch := range this.batches {
+		batch.close()
+	}
+
+	defer func() {
+		if this.locker != nil {
+			_ = this.locker.Release()
+		}
+	}()
+
+	// print log
+	/**if len(this.dsn) > 0 {
+		u, _ := url.Parse(this.dsn)
+		if u != nil && len(u.Path) > 0 {
+			remotelogs.Debug("DB", "close '"+u.Path)
+		}
+	}**/
+
+	return this.rawDB.Close()
+}
+
+func (this *DB) BeginUpdating() bool {
+	this.statusLocker.Lock()
+	defer this.statusLocker.Unlock()
+
+	if this.isClosing {
+		return false
+	}
+
+	this.countUpdating++
+	return true
+}
+
+func (this *DB) EndUpdating() {
+	this.statusLocker.Lock()
+	this.countUpdating--
+	this.statusLocker.Unlock()
+}
+
+func (this *DB) RawDB() *sql.DB {
+	return this.rawDB
+}

EdgeDNS/internal/utils/dbs/db_test.go

@@ -0,0 +1,17 @@
+// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package dbs_test
+
+import (
+	"net/url"
+	"testing"
+)
+
+func TestParseDSN(t *testing.T) {
+	var dsn = "file:/home/cache/p43/.indexes/db-3.db?cache=private&mode=ro&_journal_mode=WAL&_sync=OFF&_cache_size=88000"
+	u, err := url.Parse(dsn)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Log(u.Path) // expect: :/home/cache/p43/.indexes/db-3.db
+}

EdgeDNS/internal/utils/dbs/query_label.go

@@ -0,0 +1,24 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package dbs
+
+import "time"
+
+type QueryLabel struct {
+	manager *QueryStatManager
+	query   string
+	before  time.Time
+}
+
+func NewQueryLabel(manager *QueryStatManager, query string) *QueryLabel {
+	return &QueryLabel{
+		manager: manager,
+		query:   query,
+		before:  time.Now(),
+	}
+}
+
+func (this *QueryLabel) End() {
+	var cost = time.Since(this.before).Seconds()
+	this.manager.AddCost(this.query, cost)
+}

EdgeDNS/internal/utils/dbs/query_stat.go

@@ -0,0 +1,30 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package dbs
+
+type QueryStat struct {
+	Query   string
+	CostMin float64
+	CostMax float64
+
+	CostTotal float64
+	Calls     int64
+}
+
+func NewQueryStat(query string) *QueryStat {
+	return &QueryStat{
+		Query: query,
+	}
+}
+
+func (this *QueryStat) AddCost(cost float64) {
+	if this.CostMin == 0 || this.CostMin > cost {
+		this.CostMin = cost
+	}
+	if this.CostMax == 0 || this.CostMax < cost {
+		this.CostMax = cost
+	}
+
+	this.CostTotal += cost
+	this.Calls++
+}

EdgeDNS/internal/utils/dbs/query_stat_manager.go

@@ -0,0 +1,89 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package dbs
+
+import (
+	"fmt"
+	teaconst "github.com/TeaOSLab/EdgeDNS/internal/const"
+	"github.com/TeaOSLab/EdgeDNS/internal/events"
+	"github.com/TeaOSLab/EdgeDNS/internal/goman"
+	"github.com/iwind/TeaGo/logs"
+	"sort"
+	"strings"
+	"sync"
+	"time"
+)
+
+func init() {
+	if !teaconst.IsMain {
+		return
+	}
+
+	var ticker = time.NewTicker(5 * time.Second)
+
+	events.On(events.EventLoaded, func() {
+		if teaconst.EnableDBStat {
+			goman.New(func() {
+				for range ticker.C {
+					var stats = []string{}
+					for _, stat := range SharedQueryStatManager.TopN(10) {
+						var avg = stat.CostTotal / float64(stat.Calls)
+						var query = stat.Query
+						if len(query) > 128 {
+							query = query[:128]
+						}
+						stats = append(stats, fmt.Sprintf("%.2fms/%.2fms/%.2fms - %d - %s", stat.CostMin*1000, stat.CostMax*1000, avg*1000, stat.Calls, query))
+					}
+					logs.Println("\n========== DB STATS ==========\n" + strings.Join(stats, "\n") + "\n=============================")
+				}
+			})
+		}
+	})
+}
+
+var SharedQueryStatManager = NewQueryStatManager()
+
+type QueryStatManager struct {
+	statsMap map[string]*QueryStat // query => *QueryStat
+	locker   sync.Mutex
+}
+
+func NewQueryStatManager() *QueryStatManager {
+	return &QueryStatManager{
+		statsMap: map[string]*QueryStat{},
+	}
+}
+
+func (this *QueryStatManager) AddQuery(query string) *QueryLabel {
+	return NewQueryLabel(this, query)
+}
+
+func (this *QueryStatManager) AddCost(query string, cost float64) {
+	this.locker.Lock()
+	defer this.locker.Unlock()
+
+	stat, ok := this.statsMap[query]
+	if !ok {
+		stat = NewQueryStat(query)
+		this.statsMap[query] = stat
+	}
+	stat.AddCost(cost)
+}
+
+func (this *QueryStatManager) TopN(n int) []*QueryStat {
+	this.locker.Lock()
+	defer this.locker.Unlock()
+
+	var stats = []*QueryStat{}
+	for _, stat := range this.statsMap {
+		stats = append(stats, stat)
+	}
+	sort.Slice(stats, func(i, j int) bool {
+		return stats[i].CostMax > stats[j].CostMax
+	})
+
+	if len(stats) > n {
+		return stats[:n]
+	}
+	return stats
+}

EdgeDNS/internal/utils/dbs/query_stat_manager_test.go

@@ -0,0 +1,24 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package dbs_test
+
+import (
+	"github.com/TeaOSLab/EdgeDNS/internal/utils/dbs"
+	"github.com/iwind/TeaGo/logs"
+	"testing"
+	"time"
+)
+
+func TestQueryStatManager(t *testing.T) {
+	var manager = dbs.NewQueryStatManager()
+	{
+		var label = manager.AddQuery("sql 1")
+		time.Sleep(1 * time.Second)
+		label.End()
+	}
+	manager.AddQuery("sql 1").End()
+	manager.AddQuery("sql 2").End()
+	for _, stat := range manager.TopN(10) {
+		logs.PrintAsJSON(stat, t)
+	}
+}

EdgeDNS/internal/utils/dbs/stmt.go

@@ -0,0 +1,97 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package dbs
+
+import (
+	"context"
+	"database/sql"
+)
+
+type Stmt struct {
+	db      *DB
+	rawStmt *sql.Stmt
+	query   string
+
+	enableStat bool
+}
+
+func NewStmt(db *DB, rawStmt *sql.Stmt, query string) *Stmt {
+	return &Stmt{
+		db:      db,
+		rawStmt: rawStmt,
+		query:   query,
+	}
+}
+
+func (this *Stmt) EnableStat() {
+	this.enableStat = true
+}
+
+func (this *Stmt) ExecContext(ctx context.Context, args ...any) (sql.Result, error) {
+	// check database status
+	if this.db.BeginUpdating() {
+		defer this.db.EndUpdating()
+	} else {
+		return nil, errDBIsClosed
+	}
+
+	if this.enableStat {
+		defer SharedQueryStatManager.AddQuery(this.query).End()
+	}
+	return this.rawStmt.ExecContext(ctx, args...)
+}
+
+func (this *Stmt) Exec(args ...any) (sql.Result, error) {
+	// check database status
+	if this.db.BeginUpdating() {
+		defer this.db.EndUpdating()
+	} else {
+		return nil, errDBIsClosed
+	}
+
+	if this.enableStat {
+		defer SharedQueryStatManager.AddQuery(this.query).End()
+	}
+	return this.rawStmt.Exec(args...)
+}
+
+func (this *Stmt) QueryContext(ctx context.Context, args ...any) (*sql.Rows, error) {
+	if this.enableStat {
+		defer SharedQueryStatManager.AddQuery(this.query).End()
+	}
+	return this.rawStmt.QueryContext(ctx, args...)
+}
+
+func (this *Stmt) Query(args ...any) (*sql.Rows, error) {
+	if this.enableStat {
+		defer SharedQueryStatManager.AddQuery(this.query).End()
+	}
+	rows, err := this.rawStmt.Query(args...)
+	if err != nil {
+		return nil, err
+	}
+	var rowsErr = rows.Err()
+	if rowsErr != nil {
+		_ = rows.Close()
+		return nil, rowsErr
+	}
+	return rows, nil
+}
+
+func (this *Stmt) QueryRowContext(ctx context.Context, args ...any) *sql.Row {
+	if this.enableStat {
+		defer SharedQueryStatManager.AddQuery(this.query).End()
+	}
+	return this.rawStmt.QueryRowContext(ctx, args...)
+}
+
+func (this *Stmt) QueryRow(args ...any) *sql.Row {
+	if this.enableStat {
+		defer SharedQueryStatManager.AddQuery(this.query).End()
+	}
+	return this.rawStmt.QueryRow(args...)
+}
+
+func (this *Stmt) Close() error {
+	return this.rawStmt.Close()
+}

EdgeDNS/internal/utils/dbs/utils.go

@@ -0,0 +1,7 @@
+// Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package dbs
+
+func IsClosedErr(err error) bool {
+	return err == errDBIsClosed
+}