如何強制 IAsyncEnumerable 尊重 CancellationToken

Question

我有一個異步迭代器方法，它產生一個IAsyncEnumerable<int> （一個 stream 數字），每 200 毫秒一個數字。 此方法的調用者消耗了 stream，但希望在 1000 毫秒后停止枚舉。 因此使用了CancellationTokenSource ，並將令牌作為參數傳遞給WithCancellation擴展方法。 但是令牌不受尊重。 枚舉一直持續到所有數字都被消耗完：

static async IAsyncEnumerable<int> GetSequence()
{
    for (int i = 1; i <= 10; i++)
    {
        await Task.Delay(200);
        yield return i;
    }
}

var cts = new CancellationTokenSource(1000);
await foreach (var i in GetSequence().WithCancellation(cts.Token))
{
    Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} > {i}");
}

Output：

12:55:17.506 > 1
12:55:17.739 > 2
12:55:17.941 > 3
12:55:18.155 > 4
12:55:18.367 > 5
12:55:18.570 > 6
12:55:18.772 > 7
12:55:18.973 > 8
12:55:19.174 > 9
12:55:19.376 > 10

預期的 output 是在第 5 號之后發生的TaskCanceledException 。看來我誤解了WithCancellation實際上在做什么。 該方法只是將提供的令牌傳遞給迭代器方法，如果該方法接受一個。 否則，就像我的示例中的GetSequence()方法一樣，令牌將被忽略。 我想我的解決方案是手動查詢枚舉體內的令牌：

var cts = new CancellationTokenSource(1000);
await foreach (var i in GetSequence())
{
    cts.Token.ThrowIfCancellationRequested();
    Console.WriteLine($"{DateTime.Now:HH:mm:ss.fff} > {i}");
}

這很簡單並且效果很好。 但無論如何，我想知道是否有可能創建一個擴展方法來完成我期望WithCancellation所做的事情，以在隨后的枚舉中烘焙令牌。 這是所需方法的簽名：

public static IAsyncEnumerable<T> WithEnforcedCancellation<T>(
    this IAsyncEnumerable<T> source, CancellationToken cancellationToken)
{
    // Is it possible?
}

Answer 1

IAsyncEnumerable使用EnumeratorCancellation屬性顯式提供此機制：

static async IAsyncEnumerable<int> GetSequence([EnumeratorCancellation] CancellationToken ct = default) {
    for (int i = 1; i <= 10; i++) {
        ct.ThrowIfCancellationRequested();
        await Task.Delay(200);    // or `Task.Delay(200, ct)` if this wasn't an example
        yield return i;
    }
}

事實上，如果您給方法一個CancellationToken參數，但不添加該屬性，編譯器足以發出警告。

請注意，傳遞給.WithCancellation的令牌將覆蓋傳遞給該方法的任何本地令牌。 規格有這方面的詳細信息。

當然，這仍然只有在枚舉實際上接受CancellationToken時才有效——但取消只有在合作完成時才真正有效的事實對於任何async工作都是如此。 Yeldar 的回答對於“強制”將某種取消措施“強制”到不支持它的枚舉中是有好處的，但首選的解決方案應該是修改枚舉以自行支持取消——編譯器會盡一切努力幫助你。

Answer 2

您可以將您的邏輯提取到這樣的擴展方法中：

public static async IAsyncEnumerable<T> WithEnforcedCancellation<T>(
    this IAsyncEnumerable<T> source, CancellationToken cancellationToken)
{
    if (source == null)
        throw new ArgumentNullException(nameof(source));

    cancellationToken.ThrowIfCancellationRequested();

    await foreach (var item in source)
    {
        cancellationToken.ThrowIfCancellationRequested();
        yield return item;
    }
}

Answer 3

我認為重申你不應該這樣做很重要。 讓異步方法支持取消令牌總是更好，然后取消是您所期望的那樣立即。 如果那不可能，我仍然建議在嘗試此答案之前嘗試其他答案之一。

話雖如此，如果您無法向 async 方法添加取消支持，並且您確實需要立即終止foreach ，那么您可以繞過它。

一個技巧是將Task.WhenAny與兩個 arguments 一起使用：

1. 您從IAsyncEnumerator.MoveNextAsync()獲得的任務
2. 另一個支持取消的任務

這是簡短的版本

// Start the 'await foreach' without the new syntax
// because we need access to the ValueTask returned by MoveNextAsync()
var enumerator = source.GetAsyncEnumerator(cancellationToken);

// Combine MoveNextAsync() with another Task that can be awaited indefinitely,
// until it throws OperationCanceledException
var untilCanceled = UntilCanceled(cancellationToken);
while (await await Task.WhenAny(enumerator.MoveNextAsync().AsTask(), untilCanceled))
{
    yield return enumerator.Current;
}

為了完整性，帶有ConfigureAwait(false)和DisposeAsync()的長版本，如果您在本地運行它應該可以工作。

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;

public static class AsyncStreamHelper
{
    public static async IAsyncEnumerable<T> WithEnforcedCancellation<T>(this IAsyncEnumerable<T> source, [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        if (source == null)
            throw new ArgumentNullException(nameof(source));
        cancellationToken.ThrowIfCancellationRequested();

        // Start the 'await foreach' without the new syntax
        // because we need access to the ValueTask returned by MoveNextAsync()
        var enumerator = source.GetAsyncEnumerator(cancellationToken);
        Task<bool> moveNext = null;

        // Combine MoveNextAsync() with another Task that can be awaited indefinitely,
        // until it throws OperationCanceledException
        var untilCanceled = UntilCanceled(cancellationToken);
        try
        {
            while (
                await (
                    await Task.WhenAny(
                        (
                            moveNext = enumerator.MoveNextAsync().AsTask()
                        ),
                        untilCanceled
                    ).ConfigureAwait(false)
                )
            )
            {
                yield return enumerator.Current;
            }
        }
        finally
        {
            if (moveNext != null && !moveNext.IsCompleted)
            {
                // Disable warning CS4014 "Because this call is not awaited, execution of the current method continues before the call is completed"
#pragma warning disable 4014 // This is the behavior we want!

                moveNext.ContinueWith(async _ =>
                {
                    await enumerator.DisposeAsync();
                }, TaskScheduler.Default);
#pragma warning restore 4014
            }
            else if (enumerator != null)
            {
                await enumerator.DisposeAsync();
            }
        }
    }

    private static Task<bool> UntilCanceled(CancellationToken cancellationToken)
    {
        // This is just one possible implementation... feel free to swap out for something else
        return new Task<bool>(() => true, cancellationToken);
    }
}

public class Program
{
    public static async Task Main()
    {
        var cts = new CancellationTokenSource(500);
        var stopwatch = Stopwatch.StartNew();
        try
        {
            await foreach (var i in GetSequence().WithEnforcedCancellation(cts.Token))
            {
                Console.WriteLine($"{stopwatch.Elapsed:m':'ss'.'fff} > {i}");
            }
        }
        catch (OperationCanceledException)
        {
            Console.WriteLine($"{stopwatch.Elapsed:m':'ss'.'fff} > Canceled");
        }
    }

    static async IAsyncEnumerable<int> GetSequence()
    {
        for (int i = 1; i <= 10; i++)
        {
            await Task.Delay(200);
            yield return i;
        }
    }
}

注意事項

枚舉器返回一個 ValueTask 以提高性能（使用比常規任務更少的分配），但 ValueTask 不能與Task.WhenAny()一起使用，因此使用AsTask()會通過引入分配開銷而降低性能。

只有在最近的MoveNextAsync()完成后才能釋放枚舉器。 當請求取消時，任務更有可能仍在運行。 這就是為什么我在延續任務中添加了另一個對DisposeAsync的調用。

在這種情況下，當WithEnforcedCancellation()方法退出時，枚舉器尚未釋放。 它將在枚舉被放棄后的一段時間內被處理。 如果DisposeAsync()拋出異常，異常將丟失。 它不能冒泡調用堆棧，因為沒有調用堆棧。