Initial commit: TaskPool React panel
- React frontend with route-level code splitting - Backend rebranded from Baihu to TaskPool - DB brand migration script and local compatibility
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,251 @@
|
||||
package tasks
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/engigu/taskpool/internal/constant"
|
||||
"github.com/engigu/taskpool/internal/database"
|
||||
"github.com/engigu/taskpool/internal/logger"
|
||||
"github.com/engigu/taskpool/internal/models"
|
||||
"github.com/engigu/taskpool/internal/systime"
|
||||
"github.com/engigu/taskpool/internal/utils"
|
||||
)
|
||||
|
||||
// SendStatsService 接口定义（避免循环依赖）
|
||||
type SendStatsService interface {
|
||||
IncrementStats(taskID string, status string) error
|
||||
}
|
||||
|
||||
// TaskLogService 任务日志服务
|
||||
type TaskLogService struct {
|
||||
sendStatsService SendStatsService
|
||||
}
|
||||
|
||||
// NewTaskLogService 创建任务日志服务
|
||||
func NewTaskLogService(sendStatsService SendStatsService) *TaskLogService {
|
||||
return &TaskLogService{
|
||||
sendStatsService: sendStatsService,
|
||||
}
|
||||
}
|
||||
|
||||
// CleanConfig 清理配置
|
||||
type CleanConfig struct {
|
||||
Type string `json:"type"` // day 或 count
|
||||
Keep int `json:"keep"` // 保留天数或条数
|
||||
}
|
||||
|
||||
// CreateEmptyLog 创建一个空的日志记录（任务开始时调用）
|
||||
func (s *TaskLogService) CreateEmptyLog(taskID string, command string) (*models.TaskLog, error) {
|
||||
startTime := models.Now()
|
||||
taskLog := &models.TaskLog{
|
||||
ID: utils.GenerateID(),
|
||||
TaskID: taskID,
|
||||
Command: models.BigText(command),
|
||||
Status: "running",
|
||||
StartTime: &startTime,
|
||||
CreatedAt: models.Now(),
|
||||
}
|
||||
if err := database.DB.Create(taskLog).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 任务开始时即更新任务的 last_run 为启动时间
|
||||
database.DB.Model(&models.Task{}).Where("id = ?", taskID).Update("last_run", startTime)
|
||||
|
||||
return taskLog, nil
|
||||
}
|
||||
|
||||
// SaveTaskLog 保存或更新任务日志
|
||||
func (s *TaskLogService) SaveTaskLog(taskLog *models.TaskLog) error {
|
||||
var err error
|
||||
if taskLog.ID != "" {
|
||||
// 先检查记录是否存在，如果不存在则创建，存在则更新
|
||||
var count int64
|
||||
database.DB.Model(&models.TaskLog{}).Where("id = ?", taskLog.ID).Count(&count)
|
||||
if count > 0 {
|
||||
err = database.DB.Model(taskLog).Where("id = ?", taskLog.ID).Updates(taskLog).Error
|
||||
} else {
|
||||
err = database.DB.Create(taskLog).Error
|
||||
}
|
||||
} else {
|
||||
taskLog.ID = utils.GenerateID()
|
||||
if taskLog.CreatedAt.Time().IsZero() {
|
||||
taskLog.CreatedAt = models.Now()
|
||||
}
|
||||
err = database.DB.Create(taskLog).Error
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 更新任务的 last_run
|
||||
// 更新任务的 last_run，优先使用日志记录的启动时间
|
||||
lastRun := models.Now()
|
||||
if taskLog.StartTime != nil {
|
||||
lastRun = *taskLog.StartTime
|
||||
}
|
||||
database.DB.Model(&models.Task{}).Where("id = ?", taskLog.TaskID).Update("last_run", lastRun)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateTaskDuration 更新任务耗时（心跳）
|
||||
func (s *TaskLogService) UpdateTaskDuration(logID string, duration int64) error {
|
||||
return database.DB.Model(&models.TaskLog{}).Where("id = ?", logID).Update("duration", duration).Error
|
||||
}
|
||||
|
||||
// UpdateLogCommand 更新日志中的命令内容（用于动态生成的命令脱敏）
|
||||
func (s *TaskLogService) UpdateLogCommand(logID string, command string) error {
|
||||
return database.DB.Model(&models.TaskLog{}).Where("id = ?", logID).Update("command", models.BigText(command)).Error
|
||||
}
|
||||
|
||||
// UpdateTaskStats 更新任务统计
|
||||
func (s *TaskLogService) UpdateTaskStats(taskID string, status string) {
|
||||
if s.sendStatsService == nil {
|
||||
logger.Error("[TaskLog] SendStatsService 未初始化")
|
||||
return
|
||||
}
|
||||
err := s.sendStatsService.IncrementStats(taskID, status)
|
||||
if err != nil {
|
||||
logger.Errorf("UpdateTaskStats err: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// CleanTaskLogs 清理任务日志
|
||||
func (s *TaskLogService) CleanTaskLogs(taskID string) {
|
||||
var task models.Task
|
||||
res := database.DB.Where("id = ?", taskID).Limit(1).Find(&task)
|
||||
if res.Error != nil || res.RowsAffected == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
if task.CleanConfig == "" {
|
||||
return
|
||||
}
|
||||
|
||||
var config CleanConfig
|
||||
if err := json.Unmarshal([]byte(task.CleanConfig), &config); err != nil {
|
||||
logger.Errorf("[TaskLog] 解析清理配置失败: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
if config.Keep <= 0 {
|
||||
return
|
||||
}
|
||||
|
||||
var deleted int64
|
||||
switch config.Type {
|
||||
case "day":
|
||||
cutoff := systime.InCST(time.Now()).AddDate(0, 0, -config.Keep)
|
||||
result := database.DB.Where("task_id = ? AND created_at < ?", taskID, cutoff).Delete(&models.TaskLog{})
|
||||
deleted = result.RowsAffected
|
||||
case "count":
|
||||
var boundaryLog models.TaskLog
|
||||
res := database.DB.Where("task_id = ?", taskID).Order("id DESC").Offset(config.Keep - 1).Limit(1).Find(&boundaryLog)
|
||||
if res.Error == nil && res.RowsAffected > 0 {
|
||||
result := database.DB.Where("task_id = ? AND id < ?", taskID, boundaryLog.ID).Delete(&models.TaskLog{})
|
||||
deleted = result.RowsAffected
|
||||
}
|
||||
}
|
||||
|
||||
if deleted > 0 {
|
||||
logger.Infof("[TaskLog] 清理旧日志: #%s 共 %d 条", taskID, deleted)
|
||||
}
|
||||
}
|
||||
|
||||
// ProcessTaskCompletion 处理任务完成后的所有操作（保存日志、更新统计、清理旧日志）
|
||||
func (s *TaskLogService) ProcessTaskCompletion(taskLog *models.TaskLog) error {
|
||||
// 1. 保存/更新日志
|
||||
if err := s.SaveTaskLog(taskLog); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 2. 更新统计
|
||||
s.UpdateTaskStats(taskLog.TaskID, taskLog.Status)
|
||||
|
||||
// 3. 异步清理旧日志
|
||||
go s.CleanTaskLogs(taskLog.TaskID)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// CreateTaskLogFromAgentResult 从 Agent 结果创建任务日志
|
||||
func (s *TaskLogService) CreateTaskLogFromAgentResult(result *models.AgentTaskResult) (*models.TaskLog, error) {
|
||||
// 裁剪并压缩输出
|
||||
trimmedOutput := utils.TrimLog(result.Output, constant.MaxLogSize)
|
||||
compressed, err := utils.CompressToBase64(trimmedOutput)
|
||||
if err != nil {
|
||||
logger.Errorf("[TaskLog] 压缩日志失败: %v", err)
|
||||
compressed = ""
|
||||
}
|
||||
|
||||
logID := result.LogID
|
||||
if logID == "" {
|
||||
logID = utils.GenerateID()
|
||||
}
|
||||
|
||||
taskLog := &models.TaskLog{
|
||||
ID: logID,
|
||||
TaskID: result.TaskID,
|
||||
AgentID: &result.AgentID,
|
||||
Command: models.BigText(result.Command),
|
||||
Output: models.BigText(compressed),
|
||||
Error: models.BigText(result.Error),
|
||||
Status: result.Status,
|
||||
Duration: result.Duration,
|
||||
ExitCode: result.ExitCode,
|
||||
CreatedAt: models.Now(),
|
||||
}
|
||||
|
||||
// 处理开始和结束时间
|
||||
if result.StartTime > 0 {
|
||||
startTime := models.LocalTime(time.Unix(result.StartTime, 0))
|
||||
taskLog.StartTime = &startTime
|
||||
}
|
||||
if result.EndTime > 0 {
|
||||
endTime := models.LocalTime(time.Unix(result.EndTime, 0))
|
||||
taskLog.EndTime = &endTime
|
||||
}
|
||||
|
||||
return taskLog, nil
|
||||
}
|
||||
|
||||
// CreateTaskLogFromLocalExecution 从本地执行结果创建任务日志
|
||||
func (s *TaskLogService) CreateTaskLogFromLocalExecution(taskID string, command, output, systemErr, status string, duration int64, exitCode int, start, end time.Time, isCompressed bool) (*models.TaskLog, error) {
|
||||
var compressed string
|
||||
var err error
|
||||
|
||||
if isCompressed {
|
||||
compressed = output
|
||||
} else {
|
||||
// 裁剪并压缩输出
|
||||
trimmedOutput := utils.TrimLog(output, constant.MaxLogSize)
|
||||
compressed, err = utils.CompressToBase64(trimmedOutput)
|
||||
if err != nil {
|
||||
logger.Errorf("[TaskLog] 压缩日志失败: %v", err)
|
||||
compressed = ""
|
||||
}
|
||||
}
|
||||
|
||||
startTime := models.LocalTime(start)
|
||||
endTime := models.LocalTime(end)
|
||||
|
||||
taskLog := &models.TaskLog{
|
||||
ID: utils.GenerateID(),
|
||||
TaskID: taskID,
|
||||
Command: models.BigText(command),
|
||||
Output: models.BigText(compressed),
|
||||
Error: models.BigText(systemErr),
|
||||
Status: status,
|
||||
Duration: duration,
|
||||
ExitCode: exitCode,
|
||||
StartTime: &startTime,
|
||||
EndTime: &endTime,
|
||||
CreatedAt: models.Now(),
|
||||
}
|
||||
|
||||
return taskLog, nil
|
||||
}
|
||||
@@ -0,0 +1,287 @@
|
||||
package tasks
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/engigu/taskpool/internal/constant"
|
||||
"github.com/engigu/taskpool/internal/database"
|
||||
"github.com/engigu/taskpool/internal/models"
|
||||
"github.com/engigu/taskpool/internal/services/relation"
|
||||
"github.com/engigu/taskpool/internal/utils"
|
||||
)
|
||||
|
||||
// TaskParam 任务创建与更新参数传输对象
|
||||
type TaskParam struct {
|
||||
Name string
|
||||
Remark string
|
||||
Command string
|
||||
PreCommand string
|
||||
PostCommand string
|
||||
Tags string
|
||||
Type string
|
||||
Config string
|
||||
Schedule string
|
||||
Timeout int
|
||||
WorkDir string
|
||||
CleanConfig string
|
||||
Envs string
|
||||
Languages models.TaskLanguages
|
||||
AgentID *string
|
||||
TriggerType string
|
||||
RetryCount int
|
||||
RetryInterval int
|
||||
RandomRange int
|
||||
SourceID string
|
||||
PinType string
|
||||
Enabled bool
|
||||
}
|
||||
|
||||
type TaskService struct {
|
||||
}
|
||||
|
||||
func NewTaskService() *TaskService {
|
||||
return &TaskService{}
|
||||
}
|
||||
|
||||
func (ts *TaskService) GetTaskBySourceID(sourceID string) *models.Task {
|
||||
var task models.Task
|
||||
res := database.DB.Where("source_id = ?", sourceID).Limit(1).Find(&task)
|
||||
if res.Error != nil || res.RowsAffected == 0 {
|
||||
return nil
|
||||
}
|
||||
ts.loadTagsAndEnvs([]models.Task{task})
|
||||
|
||||
return &task
|
||||
}
|
||||
|
||||
func (ts *TaskService) CreateTask(p *TaskParam) *models.Task {
|
||||
if p.Type == "" {
|
||||
p.Type = "task"
|
||||
}
|
||||
if p.TriggerType == "" {
|
||||
p.TriggerType = constant.TriggerTypeCron
|
||||
}
|
||||
if p.PinType == "" {
|
||||
p.PinType = constant.PinTypeNone
|
||||
}
|
||||
task := &models.Task{
|
||||
ID: utils.GenerateID(),
|
||||
Name: p.Name,
|
||||
Remark: p.Remark,
|
||||
Command: models.BigText(p.Command),
|
||||
PreCommand: models.BigText(p.PreCommand),
|
||||
PostCommand: models.BigText(p.PostCommand),
|
||||
PinType: p.PinType,
|
||||
Tags: p.Tags,
|
||||
Type: p.Type,
|
||||
TriggerType: p.TriggerType,
|
||||
Config: models.BigText(p.Config),
|
||||
Schedule: p.Schedule,
|
||||
Timeout: p.Timeout,
|
||||
WorkDir: p.WorkDir,
|
||||
CleanConfig: p.CleanConfig,
|
||||
Envs: models.BigText(p.Envs),
|
||||
Languages: p.Languages,
|
||||
AgentID: p.AgentID,
|
||||
Enabled: utils.BoolPtr(true),
|
||||
RetryCount: p.RetryCount,
|
||||
RetryInterval: p.RetryInterval,
|
||||
RandomRange: p.RandomRange,
|
||||
SourceID: p.SourceID,
|
||||
CreatedAt: models.Now(),
|
||||
UpdatedAt: models.Now(),
|
||||
}
|
||||
if p.TriggerType != constant.TriggerTypeCron {
|
||||
task.NextRun = nil
|
||||
}
|
||||
database.DB.Select("*").Create(task)
|
||||
relation.DataRelation.SaveTags(task.ID, constant.RelationTypeTaskTag, p.Tags)
|
||||
task.Tags = p.Tags
|
||||
relation.DataRelation.SaveRelations(task.ID, constant.RelationTypeTaskEnv, p.Envs)
|
||||
task.Envs = models.BigText(p.Envs)
|
||||
|
||||
return task
|
||||
}
|
||||
|
||||
func (ts *TaskService) GetTasks() []models.Task {
|
||||
var tasks []models.Task
|
||||
database.DB.Find(&tasks)
|
||||
ts.loadTagsAndEnvs(tasks)
|
||||
return tasks
|
||||
}
|
||||
|
||||
// GetTasksWithPagination 分页获取任务列表
|
||||
func (ts *TaskService) GetTasksWithPagination(page, pageSize int, name string, agentID *string, tags string, taskType string, sortBy string, order string) ([]models.Task, int64) {
|
||||
var tasks []models.Task
|
||||
var total int64
|
||||
|
||||
query := database.DB.Model(&models.Task{})
|
||||
if name != "" {
|
||||
query = query.Where("name LIKE ? OR remark LIKE ?", "%"+name+"%", "%"+name+"%")
|
||||
}
|
||||
|
||||
// 标签筛选 (交集或并集均可，这里保留原本的逻辑为并集，但是利用数据关联表)
|
||||
if tags != "" {
|
||||
tagList := strings.Split(tags, ",")
|
||||
var validTags []string
|
||||
for _, tag := range tagList {
|
||||
tag = strings.TrimSpace(tag)
|
||||
if tag != "" {
|
||||
validTags = append(validTags, tag)
|
||||
}
|
||||
}
|
||||
if len(validTags) > 0 {
|
||||
var storageIDs []string
|
||||
database.DB.Model(&models.DataStorage{}).Where("type = ? AND name IN ?", constant.RelationTypeTaskTag, validTags).Pluck("id", &storageIDs)
|
||||
|
||||
var taskIDs []string
|
||||
if len(storageIDs) > 0 {
|
||||
database.DB.Model(&models.DataRelation{}).Where("type = ? AND relate_id IN ?", constant.RelationTypeTaskTag, storageIDs).Pluck("data_id", &taskIDs)
|
||||
}
|
||||
|
||||
if len(taskIDs) > 0 {
|
||||
query = query.Where("id IN ?", taskIDs)
|
||||
} else {
|
||||
query = query.Where("1 = 0")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if taskType != "" && taskType != "all" {
|
||||
query = query.Where("type = ?", taskType)
|
||||
}
|
||||
if agentID != nil {
|
||||
query = query.Where("agent_id = ?", *agentID)
|
||||
}
|
||||
|
||||
sortColumn := "created_at"
|
||||
if sortBy != "" {
|
||||
switch sortBy {
|
||||
case "name", "next_run", "last_run", "created_at", "enabled":
|
||||
sortColumn = sortBy
|
||||
}
|
||||
}
|
||||
sortOrder := "DESC"
|
||||
if strings.ToUpper(order) == "ASC" {
|
||||
sortOrder = "ASC"
|
||||
}
|
||||
|
||||
query.Count(&total)
|
||||
query.Order("pin_type DESC, " + sortColumn + " " + sortOrder).Offset((page - 1) * pageSize).Limit(pageSize).Find(&tasks)
|
||||
ts.loadTagsAndEnvs(tasks)
|
||||
|
||||
return tasks, total
|
||||
}
|
||||
|
||||
func (ts *TaskService) GetTaskByID(id string) *models.Task {
|
||||
var task models.Task
|
||||
res := database.DB.Where("id = ?", id).Limit(1).Find(&task)
|
||||
if res.Error != nil || res.RowsAffected == 0 {
|
||||
return nil
|
||||
}
|
||||
tasks := []models.Task{task}
|
||||
ts.loadTagsAndEnvs(tasks)
|
||||
|
||||
return &tasks[0]
|
||||
}
|
||||
|
||||
func (ts *TaskService) UpdateTask(id string, p *TaskParam) *models.Task {
|
||||
var task models.Task
|
||||
res := database.DB.Where("id = ?", id).Limit(1).Find(&task)
|
||||
if res.Error != nil || res.RowsAffected == 0 {
|
||||
return nil
|
||||
}
|
||||
task.Name = p.Name
|
||||
task.Remark = p.Remark
|
||||
task.Command = models.BigText(p.Command)
|
||||
task.PreCommand = models.BigText(p.PreCommand)
|
||||
task.PostCommand = models.BigText(p.PostCommand)
|
||||
task.PinType = p.PinType
|
||||
task.Schedule = p.Schedule
|
||||
task.Timeout = p.Timeout
|
||||
task.WorkDir = p.WorkDir
|
||||
task.CleanConfig = p.CleanConfig
|
||||
task.Enabled = &p.Enabled
|
||||
task.AgentID = p.AgentID
|
||||
task.Languages = p.Languages
|
||||
task.Config = models.BigText(p.Config)
|
||||
task.RetryCount = p.RetryCount
|
||||
task.RetryInterval = p.RetryInterval
|
||||
task.RandomRange = p.RandomRange
|
||||
if p.Type != "" {
|
||||
task.Type = p.Type
|
||||
}
|
||||
if p.TriggerType != "" {
|
||||
task.TriggerType = p.TriggerType
|
||||
}
|
||||
if p.SourceID != "" {
|
||||
task.SourceID = p.SourceID
|
||||
}
|
||||
|
||||
database.DB.Model(&task).Select(
|
||||
"Name", "Remark", "Command", "Tags", "Schedule", "Timeout", "WorkDir",
|
||||
"CleanConfig", "Enabled", "AgentID", "Languages",
|
||||
"RetryCount", "RetryInterval", "RandomRange", "Type",
|
||||
"TriggerType", "Config", "SourceID", "PinType",
|
||||
"PreCommand", "PostCommand",
|
||||
).Updates(&task)
|
||||
|
||||
relation.DataRelation.SaveTags(task.ID, constant.RelationTypeTaskTag, p.Tags)
|
||||
task.Tags = p.Tags
|
||||
relation.DataRelation.SaveRelations(task.ID, constant.RelationTypeTaskEnv, p.Envs)
|
||||
task.Envs = models.BigText(p.Envs)
|
||||
|
||||
return &task
|
||||
}
|
||||
|
||||
func (ts *TaskService) DeleteTask(id string) bool {
|
||||
// 同时删除关联的通知推送设置
|
||||
database.DB.Where("type = ? AND data_id = ?", constant.BindingTypeTask, id).Delete(&models.NotifyBinding{})
|
||||
relation.DataRelation.CleanRelations(id, constant.RelationTypeTaskTag)
|
||||
relation.DataRelation.CleanRelations(id, constant.RelationTypeTaskEnv)
|
||||
|
||||
result := database.DB.Where("id = ?", id).Delete(&models.Task{})
|
||||
return result.RowsAffected > 0
|
||||
}
|
||||
|
||||
func (ts *TaskService) BatchDeleteTasks(ids []string) int64 {
|
||||
// 同时删除关联的通知推送设置
|
||||
database.DB.Where("type = ? AND data_id IN ?", constant.BindingTypeTask, ids).Delete(&models.NotifyBinding{})
|
||||
database.DB.Where("type = ? AND data_id IN ?", constant.RelationTypeTaskTag, ids).Delete(&models.DataRelation{})
|
||||
database.DB.Where("type = ? AND data_id IN ?", constant.RelationTypeTaskEnv, ids).Delete(&models.DataRelation{})
|
||||
|
||||
result := database.DB.Where("id IN ?", ids).Delete(&models.Task{})
|
||||
return result.RowsAffected
|
||||
}
|
||||
|
||||
// GetAllTags 获取所有任务标签
|
||||
func (ts *TaskService) GetAllTags() ([]string, error) {
|
||||
return relation.DataRelation.GetAllTags(constant.RelationTypeTaskTag)
|
||||
}
|
||||
|
||||
func (ts *TaskService) loadTagsAndEnvs(tasks []models.Task) {
|
||||
if len(tasks) == 0 {
|
||||
return
|
||||
}
|
||||
taskIDs := make([]string, len(tasks))
|
||||
for i, t := range tasks {
|
||||
taskIDs[i] = t.ID
|
||||
}
|
||||
|
||||
tagsMap := relation.DataRelation.LoadTags(taskIDs, constant.RelationTypeTaskTag)
|
||||
envsMap := relation.DataRelation.LoadRelations(taskIDs, constant.RelationTypeTaskEnv)
|
||||
|
||||
for i, t := range tasks {
|
||||
if tags, ok := tagsMap[t.ID]; ok {
|
||||
tasks[i].Tags = strings.Join(tags, ",")
|
||||
} else {
|
||||
tasks[i].Tags = ""
|
||||
}
|
||||
|
||||
if envs, ok := envsMap[t.ID]; ok {
|
||||
tasks[i].Envs = models.BigText(strings.Join(envs, ","))
|
||||
} else {
|
||||
tasks[i].Envs = models.BigText("")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,386 @@
|
||||
package tasks
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/engigu/taskpool/internal/constant"
|
||||
"github.com/engigu/taskpool/internal/logger"
|
||||
"github.com/engigu/taskpool/internal/utils"
|
||||
)
|
||||
|
||||
const (
|
||||
// maxLogBufferLen 定义了没有换行符时的最大缓冲长度 (4KB)
|
||||
maxLogBufferLen = 4096
|
||||
)
|
||||
|
||||
var (
|
||||
// globalTinyLogManager 跟踪所有活跃的 TinyLog 实例
|
||||
globalTinyLogManager = &TinyLogManager{
|
||||
logs: make(map[string]*TinyLog),
|
||||
}
|
||||
)
|
||||
|
||||
type TinyLogManager struct {
|
||||
mu sync.RWMutex
|
||||
logs map[string]*TinyLog
|
||||
}
|
||||
|
||||
func (m *TinyLogManager) Register(log *TinyLog) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.logs[log.LogID] = log
|
||||
}
|
||||
|
||||
func (m *TinyLogManager) Unregister(logID string) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
delete(m.logs, logID)
|
||||
}
|
||||
|
||||
func (m *TinyLogManager) Get(logID string) *TinyLog {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.logs[logID]
|
||||
}
|
||||
|
||||
// GetActiveLog 通过 ID 获取活跃的 TinyLog 实例
|
||||
func GetActiveLog(logID string) *TinyLog {
|
||||
return globalTinyLogManager.Get(logID)
|
||||
}
|
||||
|
||||
// TinyLog 是一个高性能、低内存占用的日志收集器
|
||||
type TinyLog struct {
|
||||
LogID string
|
||||
mu sync.RWMutex
|
||||
file *os.File
|
||||
path string
|
||||
writer *bufio.Writer
|
||||
subscribers []chan []byte
|
||||
remainder []byte // Leftover bytes from previous write (partial lines)
|
||||
masks []string // Secrets to mask
|
||||
closed bool
|
||||
}
|
||||
|
||||
// NewTinyLog 创建一个新的 TinyLog 实例（基于临时文件存储）并注册它，支持将配置的 masks 替换为 ********
|
||||
func NewTinyLog(logID string, masks []string) (*TinyLog, error) {
|
||||
f, err := os.CreateTemp("", "task_log_*.log")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
tl := &TinyLog{
|
||||
LogID: logID,
|
||||
file: f,
|
||||
path: f.Name(),
|
||||
writer: bufio.NewWriter(f),
|
||||
subscribers: make([]chan []byte, 0),
|
||||
masks: masks,
|
||||
}
|
||||
globalTinyLogManager.Register(tl)
|
||||
return tl, nil
|
||||
}
|
||||
|
||||
// Write 实现 io.Writer 接口
|
||||
func (l *TinyLog) Write(p []byte) (n int, err error) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
if l.closed {
|
||||
return 0, os.ErrClosed
|
||||
}
|
||||
|
||||
originalInputLen := len(p)
|
||||
var payload []byte
|
||||
if len(l.remainder) > 0 {
|
||||
// 为了防止 p 和 l.remainder 底层数组有重叠或不可预期的修改，这里分配新内存
|
||||
payload = make([]byte, len(l.remainder)+len(p))
|
||||
copy(payload, l.remainder)
|
||||
copy(payload[len(l.remainder):], p)
|
||||
l.remainder = nil
|
||||
} else {
|
||||
payload = p
|
||||
}
|
||||
|
||||
// 1. 寻找最后一个换行符 (\n 或 \r)
|
||||
lastLineBreak := bytes.LastIndexAny(payload, "\n\r")
|
||||
var completeBytes []byte
|
||||
var remainder []byte
|
||||
|
||||
if lastLineBreak != -1 {
|
||||
// 2. 提取出完整的行
|
||||
completeBytes = payload[:lastLineBreak+1]
|
||||
remainder = payload[lastLineBreak+1:]
|
||||
} else {
|
||||
// 3. 没有换行符，且如果长度超过最大缓冲，强制截断并输出，防止内存无限制增长
|
||||
if len(payload) > maxLogBufferLen {
|
||||
// 寻找最后一个完整的 UTF-8 字符边界，避免乱码
|
||||
lastSafe := maxLogBufferLen
|
||||
for i := maxLogBufferLen; i > 0 && i > maxLogBufferLen-4; i-- {
|
||||
if utf8.RuneStart(payload[i-1]) {
|
||||
if !utf8.FullRune(payload[i-1 : maxLogBufferLen]) {
|
||||
lastSafe = i - 1
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
completeBytes = payload[:lastSafe]
|
||||
remainder = payload[lastSafe:]
|
||||
} else {
|
||||
// 保留当前所有内容到下一轮 (必须 copy，因为 payload 底层可能是 io.Copy 的复用 buf)
|
||||
l.remainder = make([]byte, len(payload))
|
||||
copy(l.remainder, payload)
|
||||
return originalInputLen, nil
|
||||
}
|
||||
}
|
||||
|
||||
// 4. 将剩余部分保存 (必须 copy，防止后续 Read 覆盖底层数组)
|
||||
if len(remainder) > 0 {
|
||||
l.remainder = make([]byte, len(remainder))
|
||||
copy(l.remainder, remainder)
|
||||
} else {
|
||||
l.remainder = nil
|
||||
}
|
||||
|
||||
// 5. 将完整行转换为 UTF-8 并脱敏
|
||||
text := utils.MaskSecrets(utils.ToUTF8(completeBytes), l.masks)
|
||||
outData := []byte(text)
|
||||
|
||||
// 6. 输出安全部分
|
||||
_, err = l.writer.Write(outData)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
// 6. 广播给所有订阅者
|
||||
if len(l.subscribers) > 0 {
|
||||
for _, ch := range l.subscribers {
|
||||
select {
|
||||
case ch <- outData:
|
||||
default:
|
||||
// 如果订阅者处理太慢，丢弃消息以避免阻塞写入
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return originalInputLen, nil
|
||||
}
|
||||
|
||||
// WriteString 方便地写入字符串
|
||||
func (l *TinyLog) WriteString(s string) (n int, err error) {
|
||||
return l.Write([]byte(s))
|
||||
}
|
||||
|
||||
// Subscribe 返回一个实时接收日志块的通道
|
||||
func (l *TinyLog) Subscribe() chan []byte {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
ch := make(chan []byte, 100) // Buffer to handle bursts
|
||||
l.subscribers = append(l.subscribers, ch)
|
||||
return ch
|
||||
}
|
||||
|
||||
// Unsubscribe 移除订阅者
|
||||
func (l *TinyLog) Unsubscribe(ch chan []byte) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
for i, sub := range l.subscribers {
|
||||
if sub == ch {
|
||||
l.subscribers = append(l.subscribers[:i], l.subscribers[i+1:]...)
|
||||
close(ch)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Close 完成写入，关闭文件并注销实例
|
||||
func (l *TinyLog) Close() error {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
if l.closed {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 处理剩余的字节
|
||||
if len(l.remainder) > 0 {
|
||||
text := utils.MaskSecrets(utils.ToUTF8(l.remainder), l.masks)
|
||||
data := []byte(text)
|
||||
_, _ = l.writer.Write(data)
|
||||
|
||||
// 通知订阅者最后一部分内容
|
||||
for _, ch := range l.subscribers {
|
||||
select {
|
||||
case ch <- data:
|
||||
default:
|
||||
}
|
||||
}
|
||||
l.remainder = nil
|
||||
}
|
||||
|
||||
// 将缓冲区刷新到文件
|
||||
if err := l.writer.Flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 关闭所有订阅者通道
|
||||
for _, ch := range l.subscribers {
|
||||
close(ch)
|
||||
}
|
||||
l.subscribers = nil
|
||||
|
||||
l.closed = true
|
||||
globalTinyLogManager.Unregister(l.LogID)
|
||||
return l.file.Close()
|
||||
}
|
||||
|
||||
// CompressAndCleanup 读取临时文件，进行压缩处理，返回结果并删除临时文件
|
||||
func (l *TinyLog) CompressAndCleanup() (string, error) {
|
||||
// Ensure closed
|
||||
if !l.closed {
|
||||
l.Close()
|
||||
}
|
||||
|
||||
// 打开临时文件进行读取
|
||||
f, err := os.Open(l.path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer func() {
|
||||
f.Close()
|
||||
os.Remove(l.path) // Cleanup
|
||||
}()
|
||||
|
||||
// 获取文件大小
|
||||
stat, err := f.Stat()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
size := stat.Size()
|
||||
|
||||
// 如果日志极短，免去压缩和 Base64 编码，直接以 raw: 明文形式返回
|
||||
if size <= int64(utils.MinCompressSize) {
|
||||
content, err := io.ReadAll(f)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "raw:" + string(content), nil
|
||||
}
|
||||
|
||||
// 创建压缩输出缓冲区
|
||||
var buf bytes.Buffer
|
||||
b64Writer := base64.NewEncoder(base64.StdEncoding, &buf)
|
||||
|
||||
// 使用 Pool 优化压缩
|
||||
zw := utils.GetZstdWriter(b64Writer)
|
||||
defer utils.PutZstdWriter(zw)
|
||||
|
||||
maxSize := int64(constant.MaxLogSize)
|
||||
if maxSize < 1024*1024 {
|
||||
maxSize = 1024 * 1024
|
||||
}
|
||||
|
||||
var readStart int64 = 0
|
||||
if size > maxSize {
|
||||
readStart = size - maxSize
|
||||
// 写入一条截断提示
|
||||
truncatedMsg := fmt.Sprintf("\n\n[System] 日志过长，已自动截断，仅保留末尾 %d MB...\n\n", maxSize/1024/1024)
|
||||
if _, err := zw.Write([]byte(truncatedMsg)); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
if readStart > 0 {
|
||||
if _, err := f.Seek(readStart, io.SeekStart); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
// 流处理: 文件 -> Zstd -> Base64 -> 缓冲区
|
||||
if _, err := io.Copy(zw, f); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// 关闭写入器以刷新数据
|
||||
if err := zw.Close(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := b64Writer.Close(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return "zstd:" + buf.String(), nil
|
||||
}
|
||||
|
||||
// ReadLastLines 返回日志的最后 n 行
|
||||
func (l *TinyLog) ReadLastLines(n int) ([]byte, error) {
|
||||
l.mu.RLock()
|
||||
defer l.mu.RUnlock()
|
||||
|
||||
// 刷新写入器以确保磁盘上的文件是最新的
|
||||
_ = l.writer.Flush()
|
||||
|
||||
stat, err := os.Stat(l.path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
size := stat.Size()
|
||||
var limit int64 = 65536 // 预览限制：最大 64KB
|
||||
if size < limit {
|
||||
limit = size
|
||||
}
|
||||
offset := size - limit
|
||||
|
||||
data := make([]byte, limit)
|
||||
f, err := os.Open(l.path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
_, err = f.ReadAt(data, offset)
|
||||
if err != nil && err != io.EOF {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
lines := bytes.Split(data, []byte{'\n'})
|
||||
if len(lines) > n+1 {
|
||||
return bytes.Join(lines[len(lines)-n-1:], []byte{'\n'}), nil
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// GetPath 返回临时文件路径
|
||||
func (l *TinyLog) GetPath() string {
|
||||
return l.path
|
||||
}
|
||||
|
||||
// CleanupOrphanedTinyLogs 启动时清理残留的临时日志文件
|
||||
func CleanupOrphanedTinyLogs() {
|
||||
tmpDir := os.TempDir()
|
||||
files, err := os.ReadDir(tmpDir)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
count := 0
|
||||
for _, file := range files {
|
||||
if !file.IsDir() && len(file.Name()) > 9 && file.Name()[:9] == "task_log_" && filepath.Ext(file.Name()) == ".log" {
|
||||
os.Remove(filepath.Join(tmpDir, file.Name()))
|
||||
count++
|
||||
}
|
||||
}
|
||||
if count > 0 {
|
||||
logger.Infof("[System] 清理了 %d 个残留的任务日志临时文件", count)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
package tasks
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestTinyLog_UTF8Splitting(t *testing.T) {
|
||||
tl, err := NewTinyLog("test-utf8", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create TinyLog: %v", err)
|
||||
}
|
||||
defer tl.Close()
|
||||
|
||||
// "你好" in UTF-8: E4 BD A0, E5 a5 bd
|
||||
part1 := []byte{0xE4, 0xBD} // Partial "你"
|
||||
part2 := []byte{0xA0, 0xE5, 0xA5} // Rest of "你", partial "好"
|
||||
part3 := []byte{0xBD, '\n'} // Rest of "好", newline
|
||||
|
||||
_, _ = tl.Write(part1)
|
||||
if len(tl.remainder) != 2 {
|
||||
t.Errorf("Expected remainder len 2, got %d", len(tl.remainder))
|
||||
}
|
||||
|
||||
_, _ = tl.Write(part2)
|
||||
// Currently it should collect both parts but still no newline,
|
||||
// so remainder should be 5 bytes.
|
||||
if len(tl.remainder) != 5 {
|
||||
t.Errorf("Expected remainder len 5, got %d", len(tl.remainder))
|
||||
}
|
||||
|
||||
_, _ = tl.Write(part3)
|
||||
if len(tl.remainder) != 0 {
|
||||
t.Errorf("Expected remainder len 0 after newline, got %d", len(tl.remainder))
|
||||
}
|
||||
|
||||
// Read and verify
|
||||
data, err := tl.ReadLastLines(1)
|
||||
if err != nil {
|
||||
t.Fatalf("ReadLastLines failed: %v", err)
|
||||
}
|
||||
if !bytes.Contains(data, []byte("你好")) {
|
||||
t.Errorf("Expected output to contain '你好', got %q", data)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTinyLog_CarriageReturn(t *testing.T) {
|
||||
tl, err := NewTinyLog("test-cr", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create TinyLog: %v", err)
|
||||
}
|
||||
defer tl.Close()
|
||||
|
||||
input := []byte("progress: 50%\rprogress: 100%\r\n")
|
||||
_, _ = tl.Write(input)
|
||||
|
||||
data, err := tl.ReadLastLines(10)
|
||||
if err != nil {
|
||||
t.Fatalf("ReadLastLines failed: %v", err)
|
||||
}
|
||||
|
||||
// Should contain both progress lines (or at least be split correctly)
|
||||
if !bytes.Contains(data, []byte("progress: 50%")) {
|
||||
t.Errorf("Expected output to contain 'progress: 50%%', got %q", data)
|
||||
}
|
||||
if !bytes.Contains(data, []byte("progress: 100%")) {
|
||||
t.Errorf("Expected output to contain 'progress: 100%%', got %q", data)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTinyLog_LongLineCut(t *testing.T) {
|
||||
tl, err := NewTinyLog("test-long", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create TinyLog: %v", err)
|
||||
}
|
||||
defer tl.Close()
|
||||
|
||||
// Create a buffer of maxLogBufferLen-1 bytes with a multi-byte character at the maxLogBufferLen boundary
|
||||
// We want to ensure it doesn't cut in the middle of a 3-byte char.
|
||||
longData := make([]byte, maxLogBufferLen-1)
|
||||
for i := range longData {
|
||||
longData[i] = 'A'
|
||||
}
|
||||
// "你" is E4 BD A0
|
||||
longData = append(longData, 0xE4, 0xBD, 0xA0) // This starts at index maxLogBufferLen-1.
|
||||
// Index maxLogBufferLen-1: E4
|
||||
// Index maxLogBufferLen: BD
|
||||
// Index maxLogBufferLen+1: A0
|
||||
// If we cut at maxLogBufferLen, we split E4 and BD.
|
||||
|
||||
_, _ = tl.Write(longData)
|
||||
|
||||
// Since it's > maxLogBufferLen and no newline, it should trigger the cut.
|
||||
// Our logic finds the last safe boundary before maxLogBufferLen.
|
||||
// RuneStart(E4) at maxLogBufferLen-1 is true. FullRune(E4 at maxLogBufferLen-1 in payload[:maxLogBufferLen]) is false.
|
||||
// So lastSafe should be maxLogBufferLen-1.
|
||||
|
||||
// The first maxLogBufferLen-1 bytes (all 'A') should be processed.
|
||||
// The "你" should be in remainder.
|
||||
|
||||
if len(tl.remainder) != 3 {
|
||||
t.Errorf("Expected remainder len 3 (the char '你'), got %d", len(tl.remainder))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user