Jaymr/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add persistent system task log cleanup progress

Browse Source
This commit is contained in:
CaIon
2026-06-22 19:16:53 +08:00
parent 6dc4030fdf
commit a162163b48
19 changed files with 984 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/log.go
+43 -34
View File
@@ -666,21 +666,25 @@ func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelNa
return token
}
func DeleteOldLog(ctx context.Context, targetTimestamp int64, limit int) (int64, error) {
func CountOldLog(ctx context.Context, targetTimestamp int64) (int64, error) {
var total int64
if err := LOG_DB.WithContext(ctx).Model(&Log{}).Where("created_at < ?", targetTimestamp).Count(&total).Error; err != nil {
return 0, err
}
return total, nil
}
func DeleteOldLogBatch(ctx context.Context, targetTimestamp int64, limit int) (int64, error) {
if limit <= 0 {
limit = 100
}
if nil != ctx.Err() {
return 0, ctx.Err()
}
if common.UsingLogDatabase(common.DatabaseTypeClickHouse) {
if limit <= 0 {
limit = 100
}
var total int64 = 0
for {
if nil != ctx.Err() {
return total, ctx.Err()
}
var batchCount int64
if err := LOG_DB.WithContext(ctx).Raw(`
var batchCount int64
if err := LOG_DB.WithContext(ctx).Raw(`
SELECT count() FROM (
SELECT created_at, request_id
FROM logs
@@ -688,13 +692,13 @@ SELECT count() FROM (
ORDER BY created_at ASC, request_id ASC
LIMIT ?
)`, targetTimestamp, limit).Scan(&batchCount).Error; err != nil {
return total, err
}
if batchCount == 0 {
break
}
return 0, err
}
if batchCount == 0 {
return 0, nil
}
if err := LOG_DB.WithContext(ctx).Exec(`
if err := LOG_DB.WithContext(ctx).Exec(`
ALTER TABLE logs DELETE WHERE (created_at, request_id) IN (
SELECT created_at, request_id
FROM logs
@@ -702,17 +706,22 @@ ALTER TABLE logs DELETE WHERE (created_at, request_id) IN (
ORDER BY created_at ASC, request_id ASC
LIMIT ?
) SETTINGS mutations_sync = 1`, targetTimestamp, limit).Error; err != nil {
return total, err
}
total += batchCount
if batchCount < int64(limit) {
break
}
return 0, err
}
return total, nil
return batchCount, nil
}
result := LOG_DB.WithContext(ctx).Where("created_at < ?", targetTimestamp).Limit(limit).Delete(&Log{})
if nil != result.Error {
return 0, result.Error
}
return result.RowsAffected, nil
}
func DeleteOldLog(ctx context.Context, targetTimestamp int64, limit int) (int64, error) {
if limit <= 0 {
limit = 100
}
var total int64 = 0
@@ -722,14 +731,14 @@ ALTER TABLE logs DELETE WHERE (created_at, request_id) IN (
return total, ctx.Err()
}
result := LOG_DB.Where("created_at < ?", targetTimestamp).Limit(limit).Delete(&Log{})
if nil != result.Error {
return total, result.Error
rowsAffected, err := DeleteOldLogBatch(ctx, targetTimestamp, limit)
if nil != err {
return total, err
}
total += result.RowsAffected
total += rowsAffected
if result.RowsAffected < int64(limit) {
if rowsAffected < int64(limit) {
break
}
}
model/main.go
+2
View File
@@ -294,6 +294,7 @@ func migrateDB() error {
&CustomOAuthProvider{},
&UserOAuthBinding{},
&PerfMetric{},
&SystemTask{},
)
if err != nil {
return err
@@ -343,6 +344,7 @@ func migrateDBFast() error {
{&CustomOAuthProvider{}, "CustomOAuthProvider"},
{&UserOAuthBinding{}, "UserOAuthBinding"},
{&PerfMetric{}, "PerfMetric"},
{&SystemTask{}, "SystemTask"},
}
// 动态计算migration数量，确保errChan缓冲区足够大
errChan := make(chan error, len(migrations))
model/system_task.go
+275
View File
@@ -0,0 +1,275 @@
package model
import (
"errors"
"github.com/QuantumNous/new-api/common"
"gorm.io/gorm"
)
type SystemTaskStatus string
const (
SystemTaskStatusPending SystemTaskStatus = "pending"
SystemTaskStatusRunning SystemTaskStatus = "running"
SystemTaskStatusSucceeded SystemTaskStatus = "succeeded"
SystemTaskStatusFailed SystemTaskStatus = "failed"
SystemTaskTypeLogCleanup = "log_cleanup"
)
var ErrSystemTaskLockLost = errors.New("system task lock lost")
type SystemTask struct {
ID int64 `json:"id" gorm:"primary_key;AUTO_INCREMENT"`
TaskID string `json:"task_id" gorm:"type:varchar(64);uniqueIndex"`
Type string `json:"type" gorm:"type:varchar(64);index"`
Status SystemTaskStatus `json:"status" gorm:"type:varchar(32);index"`
ActiveKey *string `json:"active_key,omitempty" gorm:"type:varchar(64);uniqueIndex"`
Payload string `json:"payload" gorm:"type:text"`
State string `json:"state" gorm:"type:text"`
Result string `json:"result" gorm:"type:text"`
Error string `json:"error" gorm:"type:text"`
LockedBy string `json:"locked_by" gorm:"type:varchar(128);index"`
LockedUntil int64 `json:"locked_until" gorm:"bigint;index"`
CreatedAt int64 `json:"created_at" gorm:"bigint;index"`
UpdatedAt int64 `json:"updated_at" gorm:"bigint;index"`
}
type SystemTaskResponse struct {
ID int64 `json:"id"`
TaskID string `json:"task_id"`
Type string `json:"type"`
Status SystemTaskStatus `json:"status"`
ActiveKey string `json:"active_key,omitempty"`
Payload any `json:"payload"`
State any `json:"state"`
Result any `json:"result"`
Error string `json:"error"`
LockedBy string `json:"locked_by"`
LockedUntil int64 `json:"locked_until"`
CreatedAt int64 `json:"created_at"`
UpdatedAt int64 `json:"updated_at"`
}
func (task *SystemTask) BeforeCreate(_ *gorm.DB) error {
now := common.GetTimestamp()
if task.CreatedAt == 0 {
task.CreatedAt = now
}
if task.UpdatedAt == 0 {
task.UpdatedAt = now
}
return nil
}
func GenerateSystemTaskID() (string, error) {
key, err := common.GenerateRandomCharsKey(32)
if err != nil {
return "", err
}
return "systask_" + key, nil
}
func CreateSystemTask(taskType string, activeKey string, payload any, state any) (*SystemTask, error) {
taskID, err := GenerateSystemTaskID()
if err != nil {
return nil, err
}
payloadText, err := marshalSystemTaskJSON(payload)
if err != nil {
return nil, err
}
stateText, err := marshalSystemTaskJSON(state)
if err != nil {
return nil, err
}
task := &SystemTask{
TaskID: taskID,
Type: taskType,
Status: SystemTaskStatusPending,
Payload: payloadText,
State: stateText,
}
if activeKey != "" {
task.ActiveKey = &activeKey
}
if err := DB.Create(task).Error; err != nil {
return nil, err
}
return task, nil
}
func GetSystemTaskByTaskID(taskID string) (*SystemTask, error) {
var task SystemTask
if err := DB.Where("task_id = ?", taskID).First(&task).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, nil
}
return nil, err
}
return &task, nil
}
func GetActiveSystemTask(taskType string) (*SystemTask, error) {
var task SystemTask
err := DB.Where("type = ? AND active_key IS NOT NULL", taskType).
Where("status IN ?", activeSystemTaskStatuses()).
Order("id desc").
First(&task).Error
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, nil
}
return nil, err
}
return &task, nil
}
func FindRunnableSystemTasks(taskType string, now int64, limit int) ([]*SystemTask, error) {
var tasks []*SystemTask
if limit <= 0 {
limit = 1
}
err := DB.Where("type = ? AND status IN ? AND (locked_until = 0 OR locked_until < ?)", taskType, activeSystemTaskStatuses(), now).
Order("id asc").
Limit(limit).
Find(&tasks).Error
return tasks, err
}
func ClaimSystemTask(id int64, taskType string, runnerID string, lockUntil int64) (*SystemTask, bool, error) {
now := common.GetTimestamp()
result := DB.Model(&SystemTask{}).
Where("id = ? AND type = ? AND status IN ? AND (locked_until = 0 OR locked_until < ? OR locked_by = ?)", id, taskType, activeSystemTaskStatuses(), now, runnerID).
Updates(map[string]any{
"status": SystemTaskStatusRunning,
"locked_by": runnerID,
"locked_until": lockUntil,
"updated_at": now,
})
if result.Error != nil {
return nil, false, result.Error
}
if result.RowsAffected == 0 {
return nil, false, nil
}
var task SystemTask
if err := DB.Where("id = ?", id).First(&task).Error; err != nil {
return nil, false, err
}
return &task, true, nil
}
func UpdateSystemTaskState(taskID string, lockedBy string, state any, lockUntil int64) error {
stateText, err := marshalSystemTaskJSON(state)
if err != nil {
return err
}
result := DB.Model(&SystemTask{}).
Where("task_id = ? AND status = ? AND locked_by = ?", taskID, SystemTaskStatusRunning, lockedBy).
Updates(map[string]any{
"state": stateText,
"locked_until": lockUntil,
"updated_at": common.GetTimestamp(),
})
if result.Error != nil {
return result.Error
}
if result.RowsAffected == 0 {
return ErrSystemTaskLockLost
}
return nil
}
func FinishSystemTask(taskID string, lockedBy string, status SystemTaskStatus, resultPayload any, errorMessage string) error {
resultText, err := marshalSystemTaskJSON(resultPayload)
if err != nil {
return err
}
result := DB.Model(&SystemTask{}).
Where("task_id = ? AND status = ? AND locked_by = ?", taskID, SystemTaskStatusRunning, lockedBy).
Updates(map[string]any{
"status": status,
"active_key": nil,
"result": resultText,
"error": errorMessage,
"locked_by": "",
"locked_until": 0,
"updated_at": common.GetTimestamp(),
})
if result.Error != nil {
return result.Error
}
if result.RowsAffected == 0 {
return ErrSystemTaskLockLost
}
return nil
}
func (task *SystemTask) DecodePayload(v any) error {
return decodeSystemTaskJSONString(task.Payload, v)
}
func (task *SystemTask) DecodeState(v any) error {
return decodeSystemTaskJSONString(task.State, v)
}
func (task *SystemTask) ToResponse() SystemTaskResponse {
activeKey := ""
if task.ActiveKey != nil {
activeKey = *task.ActiveKey
}
return SystemTaskResponse{
ID: task.ID,
TaskID: task.TaskID,
Type: task.Type,
Status: task.Status,
ActiveKey: activeKey,
Payload: decodeSystemTaskJSONValue(task.Payload),
State: decodeSystemTaskJSONValue(task.State),
Result: decodeSystemTaskJSONValue(task.Result),
Error: task.Error,
LockedBy: task.LockedBy,
LockedUntil: task.LockedUntil,
CreatedAt: task.CreatedAt,
UpdatedAt: task.UpdatedAt,
}
}
func activeSystemTaskStatuses() []string {
return []string{string(SystemTaskStatusPending), string(SystemTaskStatusRunning)}
}
func marshalSystemTaskJSON(v any) (string, error) {
if v == nil {
return "", nil
}
data, err := common.Marshal(v)
if err != nil {
return "", err
}
return string(data), nil
}
func decodeSystemTaskJSONString(data string, v any) error {
if data == "" {
return nil
}
return common.UnmarshalJsonStr(data, v)
}
func decodeSystemTaskJSONValue(data string) any {
if data == "" {
return nil
}
var value any
if err := common.UnmarshalJsonStr(data, &value); err != nil {
return data
}
return value
}
model/system_task_test.go
+83
View File
@@ -0,0 +1,83 @@
package model
import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
type testSystemTaskPayload struct {
TargetTimestamp int64 `json:"target_timestamp"`
BatchSize int `json:"batch_size"`
}
type testSystemTaskState struct {
Total int64 `json:"total"`
Processed int64 `json:"processed"`
Progress int `json:"progress"`
Remaining int64 `json:"remaining"`
}
func TestSystemTaskActiveKeyIsReleasedOnFinish(t *testing.T) {
truncateTables(t)
payload := testSystemTaskPayload{TargetTimestamp: 1000, BatchSize: 100}
state := testSystemTaskState{}
task, err := CreateSystemTask(SystemTaskTypeLogCleanup, SystemTaskTypeLogCleanup, payload, state)
require.NoError(t, err)
var decodedPayload testSystemTaskPayload
require.NoError(t, task.DecodePayload(&decodedPayload))
assert.Equal(t, payload, decodedPayload)
_, err = CreateSystemTask(SystemTaskTypeLogCleanup, SystemTaskTypeLogCleanup, payload, state)
require.Error(t, err)
activeTask, err := GetActiveSystemTask(SystemTaskTypeLogCleanup)
require.NoError(t, err)
require.NotNil(t, activeTask)
assert.Equal(t, task.TaskID, activeTask.TaskID)
runnerID := "runner-a"
claimedTask, claimed, err := ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, runnerID, common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
err = FinishSystemTask(claimedTask.TaskID, runnerID, SystemTaskStatusSucceeded, map[string]int64{"deleted_count": 0}, "")
require.NoError(t, err)
activeTask, err = GetActiveSystemTask(SystemTaskTypeLogCleanup)
require.NoError(t, err)
require.Nil(t, activeTask)
_, err = CreateSystemTask(SystemTaskTypeLogCleanup, SystemTaskTypeLogCleanup, payload, state)
require.NoError(t, err)
}
func TestSystemTaskClaimRequiresExpiredLock(t *testing.T) {
truncateTables(t)
payload := testSystemTaskPayload{TargetTimestamp: 1000, BatchSize: 100}
task, err := CreateSystemTask(SystemTaskTypeLogCleanup, SystemTaskTypeLogCleanup, payload, testSystemTaskState{})
require.NoError(t, err)
claimedTask, claimed, err := ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, "runner-a", common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
_, claimed, err = ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, "runner-b", common.GetTimestamp()+60)
require.NoError(t, err)
require.False(t, claimed)
require.NoError(t, DB.Model(claimedTask).Updates(map[string]any{
"locked_until": common.GetTimestamp() - 1,
}).Error)
claimedTask, claimed, err = ClaimSystemTask(task.ID, SystemTaskTypeLogCleanup, "runner-b", common.GetTimestamp()+60)
require.NoError(t, err)
require.True(t, claimed)
assert.Equal(t, "runner-b", claimedTask.LockedBy)
}
model/task_cas_test.go
+2
View File
@@ -48,6 +48,7 @@ func TestMain(m *testing.M) {
&UserSubscription{},
&UserOAuthBinding{},
&PerfMetric{},
&SystemTask{},
); err != nil {
panic("failed to migrate: " + err.Error())
}
@@ -71,6 +72,7 @@ func truncateTables(t *testing.T) {
DB.Exec("DELETE FROM user_subscriptions")
DB.Exec("DELETE FROM user_oauth_bindings")
DB.Exec("DELETE FROM perf_metrics")
DB.Exec("DELETE FROM system_tasks")
})
}