OvertakingLegend/Assets/Script/Tool/TaskSchedulerManager.cs

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using UnityEngine;

/// <summary>
/// 多线程任务调度系统
/// </summary>
public class TaskSchedulerManager : MBSingletonManager<TaskSchedulerManager>
{
    /// <summary>
    /// 任务返回主线程的队列
    /// </summary>
    private ConcurrentQueue<Action> mainThreadQueue = new ConcurrentQueue<Action>();

    /// <summary>
    /// 用于管理多线程任务
    /// </summary>
    private Dictionary<int, Task> tasks = new Dictionary<int, Task>();

    /// <summary>
    /// 用于管理任务的 CancellationTokenSource
    /// </summary>
    private Dictionary<int, CancellationTokenSource> taskCancellationTokens = new Dictionary<int, CancellationTokenSource>();

    /// <summary>
    /// 任务ID计数器
    /// </summary>
    private int taskIdCounter = 0;

    private void Update()
    {
        // 在主线程中执行任务队列
        while (mainThreadQueue.TryDequeue(out var task))
        {
            task.Invoke();
        }

        // 清理已完成的任务
        List<int> completedTaskIds = new List<int>();
        foreach (var taskEntry in tasks)
        {
            if (taskEntry.Value.IsCompleted)
            {
                completedTaskIds.Add(taskEntry.Key);
            }
        }

        foreach (var taskId in completedTaskIds)
        {
            tasks.Remove(taskId);
            taskCancellationTokens.Remove(taskId);
        }
    }

    /// <summary>
    /// 启动一个多线程任务
    /// </summary>
    /// <param name="taskFunc">异步任务函数</param>
    /// <param name="onCompleted">任务完成后的回调</param>
    /// <returns>任务ID</returns>
    public int RunTask(Func<CancellationToken,Task> taskFunc, Action onCompleted = null)
    {
        int taskId = taskIdCounter++;
        var cts = new CancellationTokenSource();
        taskCancellationTokens[taskId] = cts;

        var task = Task.Run(async () =>
        {
            try
            {
                await taskFunc(cts.Token);
                if (onCompleted != null)
                {
                    mainThreadQueue.Enqueue(onCompleted);
                }
            }
            catch (OperationCanceledException)
            {
                Debug.Log($"任务 {taskId} 已被取消");
            }
            catch (Exception ex)
            {
                Debug.LogError($"任务 {taskId} 执行失败: {ex.Message}");
            }
        }, cts.Token);

        tasks[taskId] = task;
        return taskId;
    }

    /// <summary>
    /// 取消单个任务
    /// </summary>
    /// <param name="taskId">任务ID</param>
    public void CancelTask(int taskId)
    {
        if (taskCancellationTokens.TryGetValue(taskId, out var cts))
        {
            cts.Cancel();
            Debug.Log($"任务 {taskId} 已取消");
        }
        else
        {
            Debug.LogWarning($"未找到任务 {taskId}");
        }
    }

    /// <summary>
    /// 等待所有任务完成
    /// </summary>
    public void WaitForAllTasks()
    {
        Task.WhenAll(tasks.Values).Wait();
    }

    /// <summary>
    /// 取消所有任务
    /// </summary>
    public void CancelAllTasks()
    {
        foreach (var cts in taskCancellationTokens.Values)
        {
            cts.Cancel();
        }
        Debug.Log("所有任务已取消");
    }

    /// <summary>
    /// 清理所有任务
    /// </summary>
    public void ClearAllTasks()
    {
        CancelAllTasks();
        tasks.Clear();
        taskCancellationTokens.Clear();
        mainThreadQueue.Clear();
    }
}