argus/rear/internal/db/db.go

package db

import (
	"context"
	"fmt"
	"gorm.io/driver/mysql"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm/logger"
	"rear/internal/config"
	"sync"
	"time"

	"gorm.io/gorm"
)

// WriteTask 写任务接口
type WriteTask interface {
	Execute(db *gorm.DB) error
	GetPriority() int // 优先级，用于任务排序
	GetID() string    // 任务ID，用于去重或追踪
}

// WriteTaskResult 写任务结果
type WriteTaskResult struct {
	TaskID string
	Error  error
}

// TaskCallback 任务回调
type TaskCallback struct {
	OnSuccess func(taskID string)
	OnError   func(taskID string, err error)
}

// DatabaseManager 数据库管理器
type DatabaseManager struct {
	db     *gorm.DB
	dbType config.DatabaseType

	// 写任务队列
	writeQueue  chan WriteTaskWithCallback
	resultQueue chan WriteTaskResult
	stopCh      chan struct{}
	wg          sync.WaitGroup

	// 读写分离配置
	maxWriteWorkers int
	maxRetries      int
}

// WriteTaskWithCallback 带回调的写任务
type WriteTaskWithCallback struct {
	Task     WriteTask
	Callback TaskCallback
}

// NewDatabaseManager 创建数据库管理器
func NewDatabaseManager(db *gorm.DB, dbType config.DatabaseType) *DatabaseManager {
	maxWorkers := 1 // SQLite默认单线程
	if dbType == config.MySQL {
		maxWorkers = 10 // MySQL可以多线程
	}

	dm := &DatabaseManager{
		db:              db,
		dbType:          dbType,
		writeQueue:      make(chan WriteTaskWithCallback, 1000), // 缓冲队列
		resultQueue:     make(chan WriteTaskResult, 1000),
		stopCh:          make(chan struct{}),
		maxWriteWorkers: maxWorkers,
		maxRetries:      3,
	}

	dm.start()
	return dm
}

// start 启动数据库管理器
func (dm *DatabaseManager) start() {
	// 启动写任务处理器
	for i := 0; i < dm.maxWriteWorkers; i++ {
		dm.wg.Add(1)
		go dm.writeWorker(i)
	}

	// 启动结果处理器
	dm.wg.Add(1)
	go dm.resultProcessor()
}

// writeWorker 写任务工作协程
func (dm *DatabaseManager) writeWorker(workerID int) {
	defer dm.wg.Done()

	for {
		select {
		case taskWithCallback := <-dm.writeQueue:
			dm.executeWriteTask(taskWithCallback)
		case <-dm.stopCh:
			return
		}
	}
}

// executeWriteTask 执行写任务
func (dm *DatabaseManager) executeWriteTask(taskWithCallback WriteTaskWithCallback) {
	task := taskWithCallback.Task
	callback := taskWithCallback.Callback

	var err error
	for retry := 0; retry < dm.maxRetries; retry++ {
		// 对于SQLite，可以使用事务来优化性能
		if dm.dbType == config.SQLite {
			err = dm.db.Transaction(func(tx *gorm.DB) error {
				return task.Execute(tx)
			})
		} else {
			err = task.Execute(dm.db)
		}

		if err == nil {
			break
		}

		// 重试延迟
		if retry < dm.maxRetries-1 {
			time.Sleep(time.Millisecond * time.Duration(100*(retry+1)))
		}
	}

	// 发送结果
	result := WriteTaskResult{
		TaskID: task.GetID(),
		Error:  err,
	}

	select {
	case dm.resultQueue <- result:
	default:
		// 结果队列满了，直接执行回调
		if err != nil && callback.OnError != nil {
			callback.OnError(task.GetID(), err)
		} else if err == nil && callback.OnSuccess != nil {
			callback.OnSuccess(task.GetID())
		}
	}
}

// resultProcessor 结果处理器
func (dm *DatabaseManager) resultProcessor() {
	defer dm.wg.Done()

	for {
		select {
		case result := <-dm.resultQueue:
			// 这里可以进行结果处理，比如记录日志、监控等
			fmt.Printf("Task %s completed with error: %v\n", result.TaskID, result.Error)
		case <-dm.stopCh:
			return
		}
	}
}

// SubmitWriteTask 提交写任务
func (dm *DatabaseManager) SubmitWriteTask(task WriteTask, callback TaskCallback) error {
	taskWithCallback := WriteTaskWithCallback{
		Task:     task,
		Callback: callback,
	}

	select {
	case dm.writeQueue <- taskWithCallback:
		return nil
	default:
		return fmt.Errorf("write queue is full")
	}
}

// SubmitWriteTaskSync 同步提交写任务
func (dm *DatabaseManager) SubmitWriteTaskSync(ctx context.Context, task WriteTask) error {
	done := make(chan error, 1)

	callback := TaskCallback{
		OnSuccess: func(taskID string) {
			done <- nil
		},
		OnError: func(taskID string, err error) {
			done <- err
		},
	}

	if err := dm.SubmitWriteTask(task, callback); err != nil {
		return err
	}

	select {
	case err := <-done:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}
}

// ReadDirect 直接读取（不经过队列）
func (dm *DatabaseManager) ReadDirect() *gorm.DB {
	return dm.db
}

// GetDB 获取数据库连接（仅用于读操作）
func (dm *DatabaseManager) GetDB() *gorm.DB {
	return dm.db
}

// Stop 停止数据库管理器
func (dm *DatabaseManager) Stop() {
	close(dm.stopCh)
	dm.wg.Wait()
}

func InitDatabase() error {
	databaseConfig := getDatabaseConfig()

	var db *gorm.DB
	var err error

	switch databaseConfig.Type {
	case config.MySQL:
		dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local",
			databaseConfig.Username, databaseConfig.Password, databaseConfig.Host, databaseConfig.Port, databaseConfig.Database)
		db, err = gorm.Open(mysql.Open(dsn), &gorm.Config{
			Logger: logger.Default.LogMode(logger.Info),
		})
	case config.SQLite:
		// SQLite特殊配置
		db, err = gorm.Open(sqlite.Open(databaseConfig.Database), &gorm.Config{
			Logger: logger.Default.LogMode(logger.Info),
		})
		if err != nil {
			return err
		}

		// SQLite优化设置
		sqlDB, _ := db.DB()
		sqlDB.SetMaxIdleConns(1)    // SQLite单连接
		sqlDB.SetMaxOpenConns(1)    // SQLite单连接
		sqlDB.SetConnMaxLifetime(0) // 永不关闭

		// SQLite WAL模式优化
		db.Exec("PRAGMA journal_mode=WAL;")
		db.Exec("PRAGMA synchronous=NORMAL;")
		db.Exec("PRAGMA cache_size=1000000;")
		db.Exec("PRAGMA foreign_keys=true;")
		db.Exec("PRAGMA temp_store=memory;")

	default:
		return fmt.Errorf("unsupported database type: %s", databaseConfig.Type)
	}

	if err != nil {
		return fmt.Errorf("failed to connect database: %w", err)
	}
	DB = db
	Manger = NewDatabaseManager(db, databaseConfig.Type)
	return nil
}

func getDatabaseConfig() config.DatabaseConfig {
	// 初始化数据库
	return config.DatabaseConfig{
		Type:         config.SQLite,
		Database:     "test.db",
		MaxIdleConns: 1,
		MaxOpenConns: 1,
		MaxLifetime:  0,
	}
}

var DB *gorm.DB
var Manger *DatabaseManager

func GetDB() *gorm.DB {
	return DB
}
func GetManger() *DatabaseManager {
	return Manger
}
func IsSQLite() bool {
	cfg := getDatabaseConfig()
	return cfg.Type == config.SQLite
}