taskq/inspect.go

// Package taskq 提供基于 Redis 的异步任务队列功能
// inspect.go 文件包含统计采集器和相关数据结构
package taskq

import (
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/hibiken/asynq"
	_ "github.com/mattn/go-sqlite3"
)

// ==================== Inspector 统计采集器 ====================

// Inspector 统计采集器，独立于 HTTP 服务运行
type Inspector struct {
	inspector *asynq.Inspector
	db        *sql.DB
	closeCh   chan struct{}
	closeOnce sync.Once
	interval  time.Duration
}

// InspectorOptions 配置统计采集器的选项
type InspectorOptions struct {
	// Interval 采集间隔，默认 2 秒
	Interval time.Duration

	// DBPath SQLite 数据库文件路径，默认为 "./taskq_stats.db"
	DBPath string
}

// NewInspector 创建新的统计采集器
func NewInspector(opts InspectorOptions) (*Inspector, error) {
	if redisClient == nil {
		return nil, fmt.Errorf("taskq: redis client not initialized, call SetRedis() first")
	}

	if opts.Interval <= 0 {
		opts.Interval = 2 * time.Second
	}

	if opts.DBPath == "" {
		opts.DBPath = "./taskq_stats.db"
	}

	// 确保目录存在
	dir := filepath.Dir(opts.DBPath)
	if dir != "" && dir != "." {
		if err := os.MkdirAll(dir, 0755); err != nil {
			return nil, fmt.Errorf("taskq: failed to create directory: %v", err)
		}
	}

	// 打开 SQLite 数据库
	db, err := sql.Open("sqlite3", opts.DBPath)
	if err != nil {
		return nil, fmt.Errorf("taskq: failed to open database: %v", err)
	}

	// 初始化数据库表
	if err := initStatsDB(db); err != nil {
		db.Close()
		return nil, fmt.Errorf("taskq: failed to init database: %v", err)
	}

	ins := &Inspector{
		inspector: asynq.NewInspectorFromRedisClient(redisClient),
		db:        db,
		closeCh:   make(chan struct{}),
		interval:  opts.Interval,
	}

	// 启动后台统计采集
	go ins.startCollector()

	return ins, nil
}

// initStatsDB 初始化数据库（Prometheus 风格：单表 + 标签）
// 设计思路：
// - 单表存储所有队列的统计数据，通过 queue 列区分
// - 复合索引支持按时间和队列两个维度高效查询
// - 类似 Prometheus 的 (timestamp, labels, value) 模型
func initStatsDB(db *sql.DB) error {
	_, err := db.Exec(`
		CREATE TABLE IF NOT EXISTS metrics (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			timestamp INTEGER NOT NULL,
			queue TEXT NOT NULL,
			active INTEGER DEFAULT 0,
			pending INTEGER DEFAULT 0,
			scheduled INTEGER DEFAULT 0,
			retry INTEGER DEFAULT 0,
			archived INTEGER DEFAULT 0,
			completed INTEGER DEFAULT 0,
			succeeded INTEGER DEFAULT 0,
			failed INTEGER DEFAULT 0
		);
		-- 按队列查询：WHERE queue = ? ORDER BY timestamp
		CREATE INDEX IF NOT EXISTS idx_metrics_queue_time ON metrics(queue, timestamp DESC);
		-- 按时间查询所有队列：WHERE timestamp BETWEEN ? AND ?
		CREATE INDEX IF NOT EXISTS idx_metrics_time ON metrics(timestamp DESC);
		-- 唯一约束：同一时间同一队列只有一条记录
		CREATE UNIQUE INDEX IF NOT EXISTS idx_metrics_unique ON metrics(timestamp, queue);
	`)
	return err
}

// Close 关闭统计采集器
func (ins *Inspector) Close() error {
	ins.closeOnce.Do(func() {
		close(ins.closeCh)
	})
	if ins.db != nil {
		ins.db.Close()
	}
	return ins.inspector.Close()
}

// startCollector 启动后台统计采集任务
func (ins *Inspector) startCollector() {
	ticker := time.NewTicker(ins.interval)
	defer ticker.Stop()

	for {
		select {
		case <-ins.closeCh:
			return
		case <-ticker.C:
			ins.collectStats()
		}
	}
}

// collectStats 采集所有队列的统计数据
func (ins *Inspector) collectStats() {
	now := time.Now().Unix()

	for queueName := range queues {
		stats, err := ins.inspector.GetQueueInfo(queueName)
		if err != nil {
			continue
		}

		qs := QueueStats{
			Queue:     queueName,
			Timestamp: now,
			Active:    stats.Active,
			Pending:   stats.Pending,
			Scheduled: stats.Scheduled,
			Retry:     stats.Retry,
			Archived:  stats.Archived,
			Completed: stats.Completed,
			Succeeded: stats.Processed - stats.Failed,
			Failed:    stats.Failed,
		}

		ins.saveMetrics(qs)
	}
}

// saveMetrics 保存统计数据到 metrics 表
func (ins *Inspector) saveMetrics(stats QueueStats) error {
	if ins.db == nil {
		return nil
	}

	_, err := ins.db.Exec(`
		INSERT OR REPLACE INTO metrics (timestamp, queue, active, pending, scheduled, retry, archived, completed, succeeded, failed)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	`, stats.Timestamp, stats.Queue, stats.Active, stats.Pending, stats.Scheduled, stats.Retry, stats.Archived, stats.Completed, stats.Succeeded, stats.Failed)

	return err
}

// GetQueueInfo 获取队列信息
func (ins *Inspector) GetQueueInfo(queueName string) (*asynq.QueueInfo, error) {
	return ins.inspector.GetQueueInfo(queueName)
}

// ListActiveTasks 获取活跃任务列表
func (ins *Inspector) ListActiveTasks(queueName string, opts ...asynq.ListOption) ([]*asynq.TaskInfo, error) {
	return ins.inspector.ListActiveTasks(queueName, opts...)
}

// ListPendingTasks 获取等待任务列表
func (ins *Inspector) ListPendingTasks(queueName string, opts ...asynq.ListOption) ([]*asynq.TaskInfo, error) {
	return ins.inspector.ListPendingTasks(queueName, opts...)
}

// ListScheduledTasks 获取计划任务列表
func (ins *Inspector) ListScheduledTasks(queueName string, opts ...asynq.ListOption) ([]*asynq.TaskInfo, error) {
	return ins.inspector.ListScheduledTasks(queueName, opts...)
}

// ListRetryTasks 获取重试任务列表
func (ins *Inspector) ListRetryTasks(queueName string, opts ...asynq.ListOption) ([]*asynq.TaskInfo, error) {
	return ins.inspector.ListRetryTasks(queueName, opts...)
}

// ListArchivedTasks 获取归档任务列表
func (ins *Inspector) ListArchivedTasks(queueName string, opts ...asynq.ListOption) ([]*asynq.TaskInfo, error) {
	return ins.inspector.ListArchivedTasks(queueName, opts...)
}

// ListCompletedTasks 获取已完成任务列表
func (ins *Inspector) ListCompletedTasks(queueName string, opts ...asynq.ListOption) ([]*asynq.TaskInfo, error) {
	return ins.inspector.ListCompletedTasks(queueName, opts...)
}

// RunTask 立即运行归档任务（重试失败任务）
func (ins *Inspector) RunTask(queueName, taskID string) error {
	return ins.inspector.RunTask(queueName, taskID)
}

// PauseQueue 暂停队列
func (ins *Inspector) PauseQueue(queueName string) error {
	return ins.inspector.PauseQueue(queueName)
}

// UnpauseQueue 恢复队列
func (ins *Inspector) UnpauseQueue(queueName string) error {
	return ins.inspector.UnpauseQueue(queueName)
}

// ==================== 统计数据结构 ====================

// QueueInfo 获取每个队列的详细信息
type QueueInfo struct {
	Name        string `json:"name"`
	Priority    int    `json:"priority"`
	Size        int    `json:"size"`         // 队列中任务总数
	Active      int    `json:"active"`       // 活跃任务数
	Pending     int    `json:"pending"`      // 等待任务数
	Scheduled   int    `json:"scheduled"`    // 计划任务数
	Retry       int    `json:"retry"`        // 重试任务数
	Archived    int    `json:"archived"`     // 归档任务数
	Completed   int    `json:"completed"`    // 已完成任务数
	Processed   int    `json:"processed"`    // 累计处理数（今日）
	Failed      int    `json:"failed"`       // 累计失败数（今日）
	Paused      bool   `json:"paused"`       // 是否暂停
	MemoryUsage int64  `json:"memory_usage"` // 内存使用（字节）
	Latency     int64  `json:"latency"`      // 延迟（毫秒）
}

// QueueStats 队列统计数据点（用于存储历史数据）
type QueueStats struct {
	Timestamp int64  `json:"t"`           // Unix 时间戳（秒）
	Queue     string `json:"q,omitempty"` // 队列名称（汇总查询时为空）
	Active    int    `json:"a"`           // 活跃任务数
	Pending   int    `json:"p"`           // 等待任务数
	Scheduled int    `json:"s"`           // 计划任务数
	Retry     int    `json:"r"`           // 重试任务数
	Archived  int    `json:"ar"`          // 归档任务数
	Completed int    `json:"c"`           // 已完成任务数
	Succeeded int    `json:"su"`          // 成功数
	Failed    int    `json:"f"`           // 失败数
}

// ==================== 全局统计数据查询 ====================

var statsDB *sql.DB
var statsDBMu sync.RWMutex

// SetStatsDB 设置全局统计数据库（供 HTTPHandler 使用）
func SetStatsDB(db *sql.DB) {
	statsDBMu.Lock()
	defer statsDBMu.Unlock()
	statsDB = db
}

// StatsQuery 统计查询参数
type StatsQuery struct {
	Queue string // 队列名称，为空则查询所有队列汇总
	Start int64  // 开始时间戳（秒），0 表示不限制
	End   int64  // 结束时间戳（秒），0 表示不限制
	Limit int    // 返回数量限制，默认 500
}

// getQueueStats 获取队列历史统计数据
func getQueueStats(queueName string, limit int) ([]QueueStats, error) {
	return getQueueStatsWithQuery(StatsQuery{
		Queue: queueName,
		Limit: limit,
	})
}

// getQueueStatsWithQuery 根据查询条件获取统计数据（Prometheus 风格单表查询）
// - 按队列查询：使用 idx_metrics_queue_time 索引
// - 按时间汇总：使用 idx_metrics_time 索引 + GROUP BY
func getQueueStatsWithQuery(q StatsQuery) ([]QueueStats, error) {
	statsDBMu.RLock()
	db := statsDB
	statsDBMu.RUnlock()

	if db == nil {
		return nil, nil
	}

	if q.Limit <= 0 {
		q.Limit = 500
	}

	var args []any
	var whereClause string
	var conditions []string

	// 构建 WHERE 条件
	if q.Queue != "" {
		conditions = append(conditions, "queue = ?")
		args = append(args, q.Queue)
	}
	if q.Start > 0 {
		conditions = append(conditions, "timestamp >= ?")
		args = append(args, q.Start)
	}
	if q.End > 0 {
		conditions = append(conditions, "timestamp <= ?")
		args = append(args, q.End)
	}

	if len(conditions) > 0 {
		whereClause = "WHERE " + strings.Join(conditions, " AND ")
	}

	var query string
	if q.Queue != "" {
		// 查询单个队列
		query = fmt.Sprintf(`
			SELECT timestamp, queue, active, pending, scheduled, retry, archived, completed, succeeded, failed
			FROM metrics
			%s
			ORDER BY timestamp DESC
			LIMIT ?
		`, whereClause)
	} else {
		// 查询所有队列汇总（按时间 GROUP BY）
		query = fmt.Sprintf(`
			SELECT timestamp, '' as queue, SUM(active), SUM(pending), SUM(scheduled), SUM(retry), SUM(archived), SUM(completed), SUM(succeeded), SUM(failed)
			FROM metrics
			%s
			GROUP BY timestamp
			ORDER BY timestamp DESC
			LIMIT ?
		`, whereClause)
	}
	args = append(args, q.Limit)

	rows, err := db.Query(query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var statsList []QueueStats
	for rows.Next() {
		var s QueueStats
		if err := rows.Scan(&s.Timestamp, &s.Queue, &s.Active, &s.Pending, &s.Scheduled, &s.Retry, &s.Archived, &s.Completed, &s.Succeeded, &s.Failed); err != nil {
			continue
		}
		statsList = append(statsList, s)
	}

	// 反转顺序，使时间从早到晚
	for i, j := 0, len(statsList)-1; i < j; i, j = i+1, j-1 {
		statsList[i], statsList[j] = statsList[j], statsList[i]
	}

	return statsList, nil
}

// GetStatsDB 返回 Inspector 的数据库连接（供外部设置给 HTTPHandler）
func (ins *Inspector) GetStatsDB() *sql.DB {
	return ins.db
}