异步屏障

异步屏障

代码

AsyncBarrier

 public class AsyncBarrier
    {
        private readonly int _participantCount;
        
        private readonly Stack<Waiter> _waiters;
        public AsyncBarrier(int participants)
        {
            if (participants <= 0)
            {
                throw new ArgumentException(nameof(participants));
            }
            this._participantCount = participants;

            this._waiters = new Stack<Waiter>(participants - 1);
        }

        public Task SignalAndWait() => this.SignalAndWaitAsync(CancellationToken.None).AsTask();

        public ValueTask SignalAndWaitAsync(CancellationToken cancellationToken)
        {
            lock (this._waiters)
            {
                if (this._waiters.Count + 1 == this._participantCount)
                {
                    while (this._waiters.Count > 0)
                    {
                        var waiter = this._waiters.Pop();
                        waiter.CompletionSource.TrySetResult(default);
                        waiter.CancellationRegistration.Dispose();
                    }

                    return new ValueTask(cancellationToken.IsCancellationRequested
                        ? Task.FromCanceled(cancellationToken)
                        : Task.CompletedTask);
                }

                TaskCompletionSource<EmptyStruct> tcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
                CancellationTokenRegistration ctr;
                if (cancellationToken.CanBeCanceled)
                {
#if NET
                    ctr = cancellationToken.Register(
                        static (tcs, ct) => ((TaskCompletionSource<EmptyStruct>)tcs!).TrySetCanceled(ct), tcs);
#else
                    ctr = cancellationToken.Register(
                        static s =>
                        {
                            var t = (Tuple<TaskCompletionSource<EmptyStruct>, CancellationToken>)s!;
                            t.Item1.TrySetCanceled(t.Item2);
                        },
                        Tuple.Create(tcs, cancellationToken));
#endif
                }
                else
                {
                    ctr = default;
                }

                this._waiters.Push(new Waiter(tcs, ctr));
                return new ValueTask(tcs.Task);
            }
        }

        private readonly struct Waiter(TaskCompletionSource<EmptyStruct> completionSource, CancellationTokenRegistration cancellationRegistration)
        {
            internal readonly TaskCompletionSource<EmptyStruct> CompletionSource => completionSource;

            internal readonly CancellationTokenRegistration CancellationRegistration => cancellationRegistration;
        }
    }

EmptyStruct

   /// <summary>
    /// 一个空结构体.
    /// </summary>
    /// <remarks>
    /// 这可以在系统上节省4个字节。当泛型类型需要类型参数但完全未使用时，对象。
    /// </remarks>
    internal readonly struct EmptyStruct
    {
        /// <summary>
        /// 获取一个空结构实例.
        /// </summary>
        internal static EmptyStruct Instance => default;
    }

介绍

AsyncBarrier 类是一个用于异步任务同步的类，它允许多个异步任务在某个点相互等待，直到所有任务都到达该点后再继续执行。这种机制在需要协调多个并发任务的场景中非常有用。

以下是 AsyncBarrier 类的关键组成部分和功能解析：

---

1. 构造函数

• 功能：初始化屏障，设置参与者的数量。
• 参数：
  • participants：表示参与同步的任务数量。
• 逻辑：
  • 如果 participants 小于或等于 0，抛出 ArgumentException 异常。
  • 初始化一个栈 _waiters，用于存储等待的任务。栈的大小为 participants - 1，因为最后一个到达的任务不需要等待。

---

2. SignalAndWait 方法

• 功能：调用 SignalAndWaitAsync 方法，并提供一个默认的 CancellationToken（即不取消）。
• 返回值：返回一个 Task，表示异步操作。

---

3. SignalAndWaitAsync 方法

• 功能：核心方法，处理任务的同步逻辑。
• 参数：
  • cancellationToken：用于支持任务取消的令牌。
• 逻辑：
  • 使用 lock 确保线程安全。
  • 检查当前等待的任务数量（_waiters.Count + 1）是否等于总参与者数量（_participantCount）：
    • 如果相等：表示所有任务都已到达屏障。此时，遍历 _waiters 栈，逐个唤醒等待的任务（通过设置 TaskCompletionSource 为完成状态），并释放它们的取消注册。
    • 如果不等：当前任务需要等待。创建一个 TaskCompletionSource 对象，并将其与 CancellationTokenRegistration 一起压入 _waiters 栈中。返回一个 ValueTask，表示当前任务正在等待。
• 取消处理：
  • 如果 cancellationToken 支持取消，则注册一个回调，当取消时，将 TaskCompletionSource 设置为取消状态。
  • 如果 cancellationToken 不支持取消，则使用默认的 CancellationTokenRegistration。

---

4. Waiter 结构体

• 功能：用于存储等待任务的状态。
• 字段：
  • TaskCompletionSource<EmptyStruct> completionSource：表示任务的完成源，用于通知任务可以继续执行。
  • CancellationTokenRegistration cancellationRegistration：表示任务的取消注册，用于支持取消操作。
• 作用：将任务的状态（完成源和取消注册）封装在一起，方便管理。

---

5. CancellationToken 处理

• 功能：支持任务的取消操作。
• 逻辑：
  • 如果 cancellationToken 被取消，则通过回调将 TaskCompletionSource 设置为取消状态。
  • 在 .NET 环境中，使用 cancellationToken.Register 注册取消回调。
  • 在非 .NET 环境中，使用 Tuple 将 TaskCompletionSource 和 CancellationToken 打包传递给回调。

---

6. EmptyStruct

• 说明：代码中未定义 EmptyStruct，但可以推测它是一个空结构体，用于表示无实际意义的返回值。

---

代码的主要用途

AsyncBarrier 类的主要用途是协调多个异步任务的执行。例如：

• 在并行计算中，多个任务需要同时完成某个阶段后，才能进入下一个阶段。
• 在分布式系统中，多个节点需要同步状态后再继续执行。

---

示例场景

假设有 3 个异步任务需要同步：

1. 任务 A 调用 SignalAndWaitAsync，发现自己是第一个到达的，进入等待状态。
2. 任务 B 调用 SignalAndWaitAsync，发现自己是第二个到达的，进入等待状态。
3. 任务 C 调用 SignalAndWaitAsync，发现自己是第三个到达的（即所有任务都已到达），唤醒任务 A 和任务 B，所有任务继续执行。

---

总结

AsyncBarrier 是一个高效的异步同步原语，适用于需要协调多个异步任务的场景。它通过 TaskCompletionSource 和 CancellationToken 实现了任务的等待和取消功能，并通过栈结构管理等待的任务。