diff --git a/tms/task-manager/internal/store/sqlite.go b/tms/task-manager/internal/store/sqlite.go index 0250757eb2c87c8e48157559f9f55c590371017b..e9d8f3f673bc459158cee15fba24cf71f34e55d0 100644 --- a/tms/task-manager/internal/store/sqlite.go +++ b/tms/task-manager/internal/store/sqlite.go @@ -28,6 +28,8 @@ import ( // ErrNotFound 在查找方法找不到对应行时返回。调用方用 errors.Is 检测它。 var ErrNotFound = errors.New("store: not found") +const cleanupBatchSize = 500 + // ----------------------------------------------------------------------------- // 值类型 // ----------------------------------------------------------------------------- @@ -897,16 +899,64 @@ func (s *Store) MarkConfigPending(agentID uint64, fetchTaskID string) error { // 维护类操作 // ----------------------------------------------------------------------------- -// Cleanup 删除早于 `olderThan` 的 ops_tasks,并级联清理 task_targets, -// 最后调用 VACUUM。返回被删除的 ops_tasks 行数。 +// Cleanup 在单个事务中删除一批早于 olderThan 的 ops_tasks 及其 task_targets。 +// 返回本批删除的 ops_tasks 行数;调用方可重复调用直到返回 0。 func (s *Store) Cleanup(olderThan time.Duration) (int64, error) { cutoff := time.Now().UTC().Add(-olderThan) - res, err := s.db.Exec(`DELETE FROM ops_tasks WHERE created_at < ?`, cutoff) + tx, err := s.db.Begin() if err != nil { - return 0, err + return 0, fmt.Errorf("begin cleanup transaction: %w", err) + } + defer tx.Rollback() + + rows, err := tx.Query(`SELECT task_id FROM ops_tasks + WHERE created_at < ? ORDER BY created_at ASC, task_id ASC LIMIT ?`, cutoff, cleanupBatchSize) + if err != nil { + return 0, fmt.Errorf("select cleanup task batch: %w", err) + } + var taskIDs []string + for rows.Next() { + var taskID string + if err := rows.Scan(&taskID); err != nil { + rows.Close() + return 0, fmt.Errorf("scan cleanup task batch: %w", err) + } + taskIDs = append(taskIDs, taskID) + } + if err := rows.Close(); err != nil { + return 0, fmt.Errorf("close cleanup task batch: %w", err) + } + if err := rows.Err(); err != nil { + return 0, fmt.Errorf("read cleanup task batch: %w", err) + } + if len(taskIDs) == 0 { + if err := tx.Commit(); err != nil { + return 0, fmt.Errorf("commit empty cleanup transaction: %w", err) + } + return 0, nil + } + + placeholders := make([]string, len(taskIDs)) + args := make([]interface{}, len(taskIDs)) + for i, taskID := range taskIDs { + placeholders[i] = "?" + args[i] = taskID + } + inClause := join(placeholders, ",") + if _, err := tx.Exec(`DELETE FROM task_targets WHERE task_id IN (`+inClause+`)`, args...); err != nil { + return 0, fmt.Errorf("delete cleanup task targets: %w", err) + } + res, err := tx.Exec(`DELETE FROM ops_tasks WHERE task_id IN (`+inClause+`)`, args...) + if err != nil { + return 0, fmt.Errorf("delete cleanup tasks: %w", err) + } + n, err := res.RowsAffected() + if err != nil { + return 0, fmt.Errorf("count cleanup tasks: %w", err) + } + if err := tx.Commit(); err != nil { + return 0, fmt.Errorf("commit cleanup transaction: %w", err) } - n, _ := res.RowsAffected() - _, _ = s.db.Exec(`DELETE FROM task_targets WHERE task_id NOT IN (SELECT task_id FROM ops_tasks)`) return n, nil } diff --git a/tms/task-manager/internal/store/sqlite_test.go b/tms/task-manager/internal/store/sqlite_test.go index ab09b5d9fef084556d4db1cd5ca71b901bffc2ae..70cfc92b18ab11fd62f103a56c22d9e6dedf0356 100644 --- a/tms/task-manager/internal/store/sqlite_test.go +++ b/tms/task-manager/internal/store/sqlite_test.go @@ -5,6 +5,7 @@ package store import ( "errors" + "fmt" "path/filepath" "reflect" "testing" @@ -657,3 +658,104 @@ func TestSweepExpiredSent_NoWriteSideEffect(t *testing.T) { t.Fatalf("status mutated by pure SELECT: got %s", rows[0].Status) } } + +func TestCleanupDeletesExpiredTasksInFixedBatches(t *testing.T) { + const batchSize = 500 + + s := newTestStore(t) + now := time.Now().UTC() + tx, err := s.DB().Begin() + if err != nil { + t.Fatalf("begin seed transaction: %v", err) + } + for i := 0; i < batchSize+2; i++ { + taskID := fmt.Sprintf("expired-%03d", i) + if _, err := tx.Exec(`INSERT INTO ops_tasks + (task_id, action, created_at, expires_at, updated_at) + VALUES (?, 'ACTION_RESTART', ?, ?, ?)`, + taskID, now.Add(-48*time.Hour), now.Add(-24*time.Hour), now.Add(-48*time.Hour)); err != nil { + tx.Rollback() + t.Fatalf("insert expired task %s: %v", taskID, err) + } + if _, err := tx.Exec(`INSERT INTO task_targets + (task_id, agent_id, created_at, updated_at) VALUES (?, ?, ?, ?)`, + taskID, i+1, now.Add(-48*time.Hour), now.Add(-48*time.Hour)); err != nil { + tx.Rollback() + t.Fatalf("insert expired target %s: %v", taskID, err) + } + } + if _, err := tx.Exec(`INSERT INTO ops_tasks + (task_id, action, created_at, expires_at, updated_at) + VALUES ('fresh', 'ACTION_RESTART', ?, ?, ?)`, now, now.Add(time.Hour), now); err != nil { + tx.Rollback() + t.Fatalf("insert fresh task: %v", err) + } + if _, err := tx.Exec(`INSERT INTO task_targets + (task_id, agent_id, created_at, updated_at) VALUES ('fresh', 9999, ?, ?)`, now, now); err != nil { + tx.Rollback() + t.Fatalf("insert fresh target: %v", err) + } + if err := tx.Commit(); err != nil { + t.Fatalf("commit seed transaction: %v", err) + } + + for call, want := range []int64{batchSize, 2, 0} { + got, err := s.Cleanup(24 * time.Hour) + if err != nil { + t.Fatalf("Cleanup call %d: %v", call+1, err) + } + if got != want { + t.Fatalf("Cleanup call %d deleted %d tasks, want %d", call+1, got, want) + } + } + + var expiredTasks, expiredTargets, freshTargets int + if err := s.DB().QueryRow(`SELECT COUNT(*) FROM ops_tasks WHERE task_id LIKE 'expired-%'`).Scan(&expiredTasks); err != nil { + t.Fatalf("count expired tasks: %v", err) + } + if err := s.DB().QueryRow(`SELECT COUNT(*) FROM task_targets WHERE task_id LIKE 'expired-%'`).Scan(&expiredTargets); err != nil { + t.Fatalf("count expired targets: %v", err) + } + if err := s.DB().QueryRow(`SELECT COUNT(*) FROM task_targets WHERE task_id = 'fresh'`).Scan(&freshTargets); err != nil { + t.Fatalf("count fresh targets: %v", err) + } + if expiredTasks != 0 || expiredTargets != 0 || freshTargets != 1 { + t.Fatalf("remaining rows: expired tasks=%d targets=%d, fresh targets=%d", expiredTasks, expiredTargets, freshTargets) + } +} + +func TestCleanupRollsBackBatchWhenTargetDeleteFails(t *testing.T) { + s := newTestStore(t) + for i, taskID := range []string{"cleanup-a", "cleanup-b"} { + mustCreate(t, s, taskID, 1) + if err := s.AddTarget(taskID, uint64(i+1)); err != nil { + t.Fatalf("AddTarget(%s): %v", taskID, err) + } + } + if _, err := s.DB().Exec(`UPDATE ops_tasks SET created_at = ?`, time.Now().UTC().Add(-48*time.Hour)); err != nil { + t.Fatalf("backdate tasks: %v", err) + } + if _, err := s.DB().Exec(`CREATE TRIGGER reject_cleanup_target_delete + BEFORE DELETE ON task_targets + WHEN OLD.task_id = 'cleanup-b' + BEGIN + SELECT RAISE(FAIL, 'target delete rejected'); + END`); err != nil { + t.Fatalf("create trigger: %v", err) + } + + if _, err := s.Cleanup(24 * time.Hour); err == nil { + t.Fatal("Cleanup succeeded after target deletion failed") + } + + var tasks, targets int + if err := s.DB().QueryRow(`SELECT COUNT(*) FROM ops_tasks WHERE task_id LIKE 'cleanup-%'`).Scan(&tasks); err != nil { + t.Fatalf("count tasks: %v", err) + } + if err := s.DB().QueryRow(`SELECT COUNT(*) FROM task_targets WHERE task_id LIKE 'cleanup-%'`).Scan(&targets); err != nil { + t.Fatalf("count targets: %v", err) + } + if tasks != 2 || targets != 2 { + t.Fatalf("rows after rollback: tasks=%d targets=%d, want 2 each", tasks, targets) + } +}