feat: add persistent system task log cleanup progress

2026-06-22 19:16:53 +08:00
parent 6dc4030fdf
commit a162163b48
19 changed files with 984 additions and 51 deletions
@@ -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)
+}