JiShe.CollectBus/modules/JiShe.CollectBus.Kafka/Internal/KafkaTaskScheduler.cs

using System.Collections.Concurrent;

namespace JiShe.CollectBus.Kafka.Internal
{
    public class KafkaTaskScheduler : TaskScheduler, IDisposable
    {
        private readonly BlockingCollection<Task> _tasksCollection=new BlockingCollection<Task> ();
        private readonly List<Thread> _workerThreads;
        private readonly object _disposeLock = new object();
        private bool _isDisposed;

        /// <summary>
        /// 当前队列中的任务数
        /// </summary>
        public int QueuedTasks => _tasksCollection.Count;

        /// <summary>
        /// 当前工作线程数
        /// </summary>
        public int WorkerThreads => _workerThreads.Count;

        /// <summary>
        /// 初始化任务调度器
        /// </summary>
        public KafkaTaskScheduler()
        {
            // 默认最大并发线程数为CPU核心数
            int MaxConcurrencyLevel = Environment.ProcessorCount;
            _workerThreads = new List<Thread>(MaxConcurrencyLevel);
            for (int i = 0; i < MaxConcurrencyLevel; i++) 
            {
                var thread = new Thread(ExecuteScheduledTasks)
                {
                    IsBackground = true,
                    Name = $"KafkaWorkerTask-{i + 1}"
                };
                thread.Start();
                _workerThreads.Add(thread);
            }
          
        }

        /// <summary>
        /// 扩容工作线程调度
        /// 可以启动多个工作线程来处理任务
        /// </summary>
        /// <param name="taskNum">扩展独立线程数(默认为1)</param>
        public void WorkerThreadExpansion(int taskNum = 1)
        {
            int currCount = WorkerThreads+1;
            Parallel.For(0, taskNum, (index) =>
            {
                var thread = new Thread(ExecuteScheduledTasks)
                {
                    IsBackground = true,
                    Name = $"KafkaWorkerTask-{index+ currCount}"
                };
                thread.Start();
                _workerThreads.Add(thread);
            });
        }

        /// <summary>
        /// 工作线程执行循环
        /// </summary>
        private void ExecuteScheduledTasks()
        {
            try
            {
                foreach (var task in _tasksCollection.GetConsumingEnumerable())
                {
                    TryExecuteTaskSafely(task);
                }
            }
            catch (OperationCanceledException) { }
            catch (ObjectDisposedException) { }
        }

        /// <summary>
        /// 安全执行任务并处理异常
        /// </summary>
        private void TryExecuteTaskSafely(Task task)
        {
            try
            {
                TryExecuteTask(task);
            }
            catch (OperationCanceledException){}
            catch (Exception ex)
            {
                OnExceptionOccurred(ex);
            }
        }

        #region TaskScheduler 重写方法
        protected override IEnumerable<Task> GetScheduledTasks()
        {
            ThrowIfDisposed();
            return _tasksCollection.ToList();
        }

        protected override void QueueTask(Task task)
        {
            ThrowIfDisposed();
            _tasksCollection.Add(task);
        }

        protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)
        {
            // 禁止内联执行以强制所有任务在专用线程执行
            return false;
        }
        #endregion

        #region 释放资源
        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        protected virtual void Dispose(bool disposing)
        {
            lock (_disposeLock)
            {
                if (_isDisposed) return;

                if (disposing)
                {
                    // 停止接收新任务
                    _tasksCollection.CompleteAdding();

                    // 等待所有工作线程退出
                    foreach (var thread in _workerThreads)
                    {
                        if (thread.IsAlive)
                        {
                            thread.Join(TimeSpan.FromSeconds(5));
                        }
                    }
                    // 释放资源
                    _tasksCollection.Dispose();
                }
                _isDisposed = true;
            }
        }

        private void ThrowIfDisposed()
        {
            if (_isDisposed)
            {
                throw new ObjectDisposedException(GetType().Name);
            }
        }
        #endregion

        #region 异常事件处理
        /// <summary>
        /// 任务执行异常时触发
        /// </summary>
        public event Action<Exception>? ExceptionEvent;

        private void OnExceptionOccurred(Exception ex)
        {
            ExceptionEvent?.Invoke(ex);
        }
        #endregion
    }
}