支持异步操作的自定义线程池

Question

我想要一个满足以下要求的自定义线程池：

1. 实际线程是根据池容量预先分配的。如果需要生成并发任务，实际工作可以免费使用标准 .NET 线程池。
2. 池必须能够返回空闲线程的数量。返回的数量可能小于实际空闲线程的数量，但不能大于。当然，数字越准确越好。
3. 将工作排队到池应返回相应的Task，这应该与基于任务的 API 相得益彰。
4. 新 最大作业容量（或并行度）应该可以动态调整。尝试减少容量不必立即生效，但增加容量应该立即生效。

第一项的基本原理如下：

• 机器不应该同时运行超过 N 个工作项，其中 N 相对较小 - 在 10 到 30 之间。
• 从数据库中获取工作，如果获取了 K 个项目，那么我们要确保有 K 个空闲线程来立即开始工作。从数据库中获取工作，但仍等待下一个可用线程的情况是不可接受的。

最后一项也解释了空闲线程计数的原因——我将从数据库中获取那么多工作项。它还解释了为什么报告的空闲线程数绝不能高于实际的 - 否则我可能会获取更多可以立即启动的工作。

无论如何，这是我的实现以及一个测试它的小程序（BJE 代表后台作业引擎）：

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

namespace TaskStartLatency
{
    public class BJEThreadPool
    {
        private sealed class InternalTaskScheduler : TaskScheduler
        {
            private int m_idleThreadCount;
            private readonly BlockingCollection<Task> m_bus;

            public InternalTaskScheduler(int threadCount, BlockingCollection<Task> bus)
            {
                m_idleThreadCount = threadCount;
                m_bus = bus;
            }

            public void RunInline(Task task)
            {
                Interlocked.Decrement(ref m_idleThreadCount);
                try
                {
                    TryExecuteTask(task);
                }
                catch
                {
                    // The action is responsible itself for the error handling, for the time being...
                }
                Interlocked.Increment(ref m_idleThreadCount);
            }

            public int IdleThreadCount
            {
                get { return m_idleThreadCount; }
            }

            #region Overrides of TaskScheduler

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

            protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)
            {
                return TryExecuteTask(task);
            }

            protected override IEnumerable<Task> GetScheduledTasks()
            {
                throw new NotSupportedException();
            }

            #endregion

            public void DecrementIdleThreadCount()
            {
                Interlocked.Decrement(ref m_idleThreadCount);
            }
        }

        private class ThreadContext
        {
            private readonly InternalTaskScheduler m_ts;
            private readonly BlockingCollection<Task> m_bus;
            private readonly CancellationTokenSource m_cts;
            public readonly Thread Thread;

            public ThreadContext(string name, InternalTaskScheduler ts, BlockingCollection<Task> bus, CancellationTokenSource cts)
            {
                m_ts = ts;
                m_bus = bus;
                m_cts = cts;
                Thread = new Thread(Start)
                {
                    IsBackground = true,
                    Name = name
                };
                Thread.Start();
            }

            private void Start()
            {
                try
                {
                    foreach (var task in m_bus.GetConsumingEnumerable(m_cts.Token))
                    {
                        m_ts.RunInline(task);
                    }
                }
                catch (OperationCanceledException)
                {
                }
                m_ts.DecrementIdleThreadCount();
            }
        }

        private readonly InternalTaskScheduler m_ts;
        private readonly CancellationTokenSource m_cts = new CancellationTokenSource();
        private readonly BlockingCollection<Task> m_bus = new BlockingCollection<Task>();
        private readonly List<ThreadContext> m_threadCtxs = new List<ThreadContext>();

        public BJEThreadPool(int threadCount)
        {
            m_ts = new InternalTaskScheduler(threadCount, m_bus);
            for (int i = 0; i < threadCount; ++i)
            {
                m_threadCtxs.Add(new ThreadContext("BJE Thread " + i, m_ts, m_bus, m_cts));
            }
        }

        public void Terminate()
        {
            m_cts.Cancel();
            foreach (var t in m_threadCtxs)
            {
                t.Thread.Join();
            }
        }

        public Task Run(Action<CancellationToken> action)
        {
            return Task.Factory.StartNew(() => action(m_cts.Token), m_cts.Token, TaskCreationOptions.DenyChildAttach, m_ts);
        }
        public Task Run(Action action)
        {
            return Task.Factory.StartNew(action, m_cts.Token, TaskCreationOptions.DenyChildAttach, m_ts);
        }

        public int IdleThreadCount
        {
            get { return m_ts.IdleThreadCount; }
        }
    }

    class Program
    {
        static void Main()
        {
            const int THREAD_COUNT = 32;
            var pool = new BJEThreadPool(THREAD_COUNT);
            var tcs = new TaskCompletionSource<bool>();
            var tasks = new List<Task>();
            var allRunning = new CountdownEvent(THREAD_COUNT);

            for (int i = pool.IdleThreadCount; i > 0; --i)
            {
                var index = i;
                tasks.Add(pool.Run(cancellationToken =>
                {
                    Console.WriteLine("Started action " + index);
                    allRunning.Signal();
                    tcs.Task.Wait(cancellationToken);
                    Console.WriteLine("  Ended action " + index);
                }));
            }

            Console.WriteLine("pool.IdleThreadCount = " + pool.IdleThreadCount);

            allRunning.Wait();
            Debug.Assert(pool.IdleThreadCount == 0);

            int expectedIdleThreadCount = THREAD_COUNT;
            Console.WriteLine("Press [c]ancel, [e]rror, [a]bort or any other key");
            switch (Console.ReadKey().KeyChar)
            {
            case 'c':
                Console.WriteLine("Cancel All");
                tcs.TrySetCanceled();
                break;
            case 'e':
                Console.WriteLine("Error All");
                tcs.TrySetException(new Exception("Failed"));
                break;
            case 'a':
                Console.WriteLine("Abort All");
                pool.Terminate();
                expectedIdleThreadCount = 0;
                break;
            default:
                Console.WriteLine("Done All");
                tcs.TrySetResult(true);
                break;
            }
            try
            {
                Task.WaitAll(tasks.ToArray());
            }
            catch (AggregateException exc)
            {
                Console.WriteLine(exc.Flatten().InnerException.Message);
            }

            Debug.Assert(pool.IdleThreadCount == expectedIdleThreadCount);

            pool.Terminate();
            Console.WriteLine("Press any key");
            Console.ReadKey();
        }
    }
}

这是一个非常简单的实现，它似乎正在工作。但是，有一个问题——BJEThreadPool.Run 方法不接受异步方法。 IE。我的实现不允许我添加以下重载：

public Task Run(Func<CancellationToken, Task> action)
{
    return Task.Factory.StartNew(() => action(m_cts.Token), m_cts.Token, TaskCreationOptions.DenyChildAttach, m_ts).Unwrap();
}
public Task Run(Func<Task> action)
{
    return Task.Factory.StartNew(action, m_cts.Token, TaskCreationOptions.DenyChildAttach, m_ts).Unwrap();
}

我在InternalTaskScheduler.RunInline 中使用的模式在这种情况下不起作用。

那么，我的问题是如何添加对异步工作项的支持？只要满足文章开头概述的要求，我就可以更改整个设计。

编辑

我想澄清所需池的预期用途。请注意以下代码：

if (pool.IdleThreadCount == 0)
{
  return;
}

foreach (var jobData in FetchFromDB(pool.IdleThreadCount))
{
  pool.Run(CreateJobAction(jobData));
}

注意事项：

1. 代码将定期运行，比如每 1 分钟一次。
2. 代码将由多台机器同时运行，监视同一个数据库。
3. FetchFromDB 将使用Using SQL Server as a DB queue with multiple clients 中描述的技术从数据库中原子地获取和锁定工作。
4. CreateJobAction 将调用由jobData 表示的代码（工作代码）并在该代码完成后关闭工作。作业代码超出了我的控制范围，它几乎可以是任何东西 - 繁重的 CPU 绑定代码或轻量异步 IO 绑定代码，写得不好的同步 IO 绑定代码或所有代码的混合。它可以运行几分钟，也可以运行几个小时。关闭工作是我的代码，它将通过异步 IO 绑定代码。因此，返回的作业操作的签名是异步方法的签名。

第 2 项强调了正确识别空闲线程数量的重要性。如果有 900 个待处理的工作项和 10 台代理机器，我不能允许代理获取 300 个工作项并将它们排入线程池。为什么？因为，代理不太可能同时运行 300 个工作项。它肯定会运行一些，但其他的会在线程池工作队列中等待。假设它将运行 100 并让 200 等待（即使 100 可能遥不可及）。这使用了 3 个满载代理和 7 个空闲代理。但实际上同时处理 900 个工作项中的只有 300 个！！！

我的目标是最大限度地在可用代理之间传播工作。理想情况下，我应该评估代理的负载和待处理工作的“重量”，但这是一项艰巨的任务，并为未来的版本保留。现在，我希望为每个代理分配最大作业容量，以便在不重新启动代理的情况下提供动态增加/减少它的方法。

接下来的观察。这项工作可能需要很长时间才能运行，并且可能都是同步代码。据我了解，将线程池线程用于此类工作是不可取的。

编辑 2

有一种说法，TaskScheduler 仅适用于 CPU 密集型工作。但是如果我不知道工作的性质怎么办？我的意思是它是一个通用的后台作业引擎，它运行数千种不同的作业。我没有办法告诉“该工作受 CPU 限制”和“该工作受同步 IO 限制”，而另一个受异步 IO 限制。我希望我能，但我不能。

编辑 3

最后，我没有使用SemaphoreSlim，但我也没有使用TaskScheduler - 它最终让我的厚脸皮说这是不合适且完全错误的，而且它使代码过于复杂。

不过，我还是没看到SemaphoreSlim 是怎么回事。建议的模式：

public async Task Enqueue(Func<Task> taskGenerator)
{
    await semaphore.WaitAsync();
    try
    {
        await taskGenerator();
    }
    finally
    {
        semaphore.Release();
    }
}

期望 taskGenerator 要么是异步 IO 绑定代码，要么打开一个新线程。但是，我无法确定要执行的工作是一项还是另一项。另外，正如我从SemaphoreSlim.WaitAsync continuation code 了解到的，如果信号量被解锁，WaitAsync() 后面的代码将在同一个线程上运行，这对我来说不是很好。

无论如何，以下是我的实现，以防万一。不幸的是，我还没有理解如何动态减少池线程数，但这是另一个问题的主题。

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

namespace TaskStartLatency
{
    public interface IBJEThreadPool
    {
        void SetThreadCount(int threadCount);
        void Terminate();
        Task Run(Action action);
        Task Run(Action<CancellationToken> action);
        Task Run(Func<Task> action);
        Task Run(Func<CancellationToken, Task> action);
        int IdleThreadCount { get; }
    }

    public class BJEThreadPool : IBJEThreadPool
    {
        private interface IActionContext
        {
            Task Run(CancellationToken ct);
            TaskCompletionSource<object> TaskCompletionSource { get; }
        }

        private class ActionContext : IActionContext
        {
            private readonly Action m_action;

            public ActionContext(Action action)
            {
                m_action = action;
                TaskCompletionSource = new TaskCompletionSource<object>();
            }

            #region Implementation of IActionContext

            public Task Run(CancellationToken ct)
            {
                m_action();
                return null;
            }

            public TaskCompletionSource<object> TaskCompletionSource { get; private set; }

            #endregion
        }
        private class CancellableActionContext : IActionContext
        {
            private readonly Action<CancellationToken> m_action;

            public CancellableActionContext(Action<CancellationToken> action)
            {
                m_action = action;
                TaskCompletionSource = new TaskCompletionSource<object>();
            }

            #region Implementation of IActionContext

            public Task Run(CancellationToken ct)
            {
                m_action(ct);
                return null;
            }

            public TaskCompletionSource<object> TaskCompletionSource { get; private set; }

            #endregion
        }
        private class AsyncActionContext : IActionContext
        {
            private readonly Func<Task> m_action;

            public AsyncActionContext(Func<Task> action)
            {
                m_action = action;
                TaskCompletionSource = new TaskCompletionSource<object>();
            }

            #region Implementation of IActionContext

            public Task Run(CancellationToken ct)
            {
                return m_action();
            }

            public TaskCompletionSource<object> TaskCompletionSource { get; private set; }

            #endregion
        }
        private class AsyncCancellableActionContext : IActionContext
        {
            private readonly Func<CancellationToken, Task> m_action;

            public AsyncCancellableActionContext(Func<CancellationToken, Task> action)
            {
                m_action = action;
                TaskCompletionSource = new TaskCompletionSource<object>();
            }

            #region Implementation of IActionContext

            public Task Run(CancellationToken ct)
            {
                return m_action(ct);
            }

            public TaskCompletionSource<object> TaskCompletionSource { get; private set; }

            #endregion
        }

        private readonly CancellationTokenSource m_ctsTerminateAll = new CancellationTokenSource();
        private readonly BlockingCollection<IActionContext> m_bus = new BlockingCollection<IActionContext>();
        private readonly LinkedList<Thread> m_threads = new LinkedList<Thread>();
        private int m_idleThreadCount;

        private static int s_threadCount;

        public BJEThreadPool(int threadCount)
        {
            ReserveAdditionalThreads(threadCount);
        }

        private void ReserveAdditionalThreads(int n)
        {
            for (int i = 0; i < n; ++i)
            {
                var index = Interlocked.Increment(ref s_threadCount) - 1;

                var t = new Thread(Start)
                {
                    IsBackground = true,
                    Name = "BJE Thread " + index
                };
                Interlocked.Increment(ref m_idleThreadCount);
                t.Start();

                m_threads.AddLast(t);
            }
        }

        private void Start()
        {
            try
            {
                foreach (var actionContext in m_bus.GetConsumingEnumerable(m_ctsTerminateAll.Token))
                {
                    RunWork(actionContext).Wait();
                }
            }
            catch (OperationCanceledException)
            {
            }
            catch
            {
                // Should never happen - log the error
            }

            Interlocked.Decrement(ref m_idleThreadCount);
        }

        private async Task RunWork(IActionContext actionContext)
        {
            Interlocked.Decrement(ref m_idleThreadCount);
            try
            {
                var task = actionContext.Run(m_ctsTerminateAll.Token);
                if (task != null)
                {
                    await task;
                }
                actionContext.TaskCompletionSource.SetResult(null);
            }
            catch (OperationCanceledException)
            {
                actionContext.TaskCompletionSource.TrySetCanceled();
            }
            catch (Exception exc)
            {
                actionContext.TaskCompletionSource.TrySetException(exc);
            }
            Interlocked.Increment(ref m_idleThreadCount);
        }

        private Task PostWork(IActionContext actionContext)
        {
            m_bus.Add(actionContext);
            return actionContext.TaskCompletionSource.Task;
        }

        #region Implementation of IBJEThreadPool

        public void SetThreadCount(int threadCount)
        {
            if (threadCount > m_threads.Count)
            {
                ReserveAdditionalThreads(threadCount - m_threads.Count);
            }
            else if (threadCount < m_threads.Count)
            {
                throw new NotSupportedException();
            }
        }
        public void Terminate()
        {
            m_ctsTerminateAll.Cancel();
            foreach (var t in m_threads)
            {
                t.Join();
            }
        }

        public Task Run(Action action)
        {
            return PostWork(new ActionContext(action));
        }
        public Task Run(Action<CancellationToken> action)
        {
            return PostWork(new CancellableActionContext(action));
        }
        public Task Run(Func<Task> action)
        {
            return PostWork(new AsyncActionContext(action));
        }
        public Task Run(Func<CancellationToken, Task> action)
        {
            return PostWork(new AsyncCancellableActionContext(action));
        }

        public int IdleThreadCount
        {
            get { return m_idleThreadCount; }
        }

        #endregion
    }

    public static class Extensions
    {
        public static Task WithCancellation(this Task task, CancellationToken token)
        {
            return task.ContinueWith(t => t.GetAwaiter().GetResult(), token);
        }
    }

    class Program
    {
        static void Main()
        {
            const int THREAD_COUNT = 16;
            var pool = new BJEThreadPool(THREAD_COUNT);
            var tcs = new TaskCompletionSource<bool>();
            var tasks = new List<Task>();
            var allRunning = new CountdownEvent(THREAD_COUNT);

            for (int i = pool.IdleThreadCount; i > 0; --i)
            {
                var index = i;
                tasks.Add(pool.Run(async ct =>
                {
                    Console.WriteLine("Started action " + index);
                    allRunning.Signal();
                    await tcs.Task.WithCancellation(ct);
                    Console.WriteLine("  Ended action " + index);
                }));
            }

            Console.WriteLine("pool.IdleThreadCount = " + pool.IdleThreadCount);

            allRunning.Wait();
            Debug.Assert(pool.IdleThreadCount == 0);

            int expectedIdleThreadCount = THREAD_COUNT;
            Console.WriteLine("Press [c]ancel, [e]rror, [a]bort or any other key");
            switch (Console.ReadKey().KeyChar)
            {
            case 'c':
                Console.WriteLine("ancel All");
                tcs.TrySetCanceled();
                break;
            case 'e':
                Console.WriteLine("rror All");
                tcs.TrySetException(new Exception("Failed"));
                break;
            case 'a':
                Console.WriteLine("bort All");
                pool.Terminate();
                expectedIdleThreadCount = 0;
                break;
            default:
                Console.WriteLine("Done All");
                tcs.TrySetResult(true);
                break;
            }

            try
            {
                Task.WaitAll(tasks.ToArray());
            }
            catch (AggregateException exc)
            {
                Console.WriteLine(exc.Flatten().InnerException.Message);
            }

            Debug.Assert(pool.IdleThreadCount == expectedIdleThreadCount);

            pool.Terminate();
            Console.WriteLine("Press any key");
            Console.ReadKey();
        }
    }
}

Answer 1

异步“工作项”通常基于异步 IO。异步 IO 在运行时不使用线程。任务调度程序用于执行 CPU 工作（基于委托的任务）。 TaskScheduler 的概念不适用。您不能使用自定义 TaskScheduler 来影响异步代码的作用。

让您的工作项自行节流：

static SemaphoreSlim sem = new SemaphoreSlim(maxDegreeOfParallelism); //shared object

async Task MyWorkerFunction()
{
    await sem.WaitAsync();
    try
    {
        MyWork();
    }
    finally
    {
        sem.Release();
    }
}

Answer 2

As mentioned in another answer by usr 你不能用TaskScheduler 来做到这一点，因为这仅适用于 CPU 绑定的工作，不限制所有类型工作的并行化水平，无论是否并行。他还向您展示了如何使用SemaphoreSlim 来异步限制并行度。

您可以对此进行扩展，以几种方式概括这些概念。似乎对您最有利的是创建一种特殊类型的队列，该队列接受返回 Task 的操作，并以达到给定最大并行化程度的方式执行它们。

public class FixedParallelismQueue
{
    private SemaphoreSlim semaphore;
    public FixedParallelismQueue(int maxDegreesOfParallelism)
    {
        semaphore = new SemaphoreSlim(maxDegreesOfParallelism,
            maxDegreesOfParallelism);
    }

    public async Task<T> Enqueue<T>(Func<Task<T>> taskGenerator)
    {
        await semaphore.WaitAsync();
        try
        {
            return await taskGenerator();
        }
        finally
        {
            semaphore.Release();
        }
    }
    public async Task Enqueue(Func<Task> taskGenerator)
    {
        await semaphore.WaitAsync();
        try
        {
            await taskGenerator();
        }
        finally
        {
            semaphore.Release();
        }
    }
}

这允许您为您的应用程序创建一个具有固定并行化程度的队列（如果需要，您甚至可以有几个单独的队列）。然后，您可以提供在完成时返回 Task 的操作，队列将在可能的情况下安排它，并返回表示该工作单元何时完成的 Task。