it-roy-ru.com

Как объединить TaskCompletionSource и CancellationTokenSource?

У меня есть такой код (здесь упрощенно), который ждет завершения задачи:

var task_completion_source = new TaskCompletionSource<bool>();
observable.Subscribe(b => 
   { 
      if (b) 
          task_completion_source.SetResult(true); 
   });
await task_completion_source.Task;    

Идея состоит в том, чтобы подписаться и дождаться true в потоке логических значений. Это завершает «задачу», и я могу выйти за пределы await.

Однако я бы хотел отменить - но не подписку, а ожидание. Я хотел бы передать токен отмены (каким-либо образом) в task_completion_source, поэтому, когда я отменяю источник токена, await будет двигаться дальше.

Как это сделать?

Обновление : CancellationTokenSource является внешним по отношению к этому коду, все, что у меня есть, это токен из него.

11
astrowalker

Если я вас правильно понимаю, вы можете сделать это так:

using (ct.Register(() => {
    // this callback will be executed when token is cancelled
    task_comletion_source.TrySetCanceled();
})) {
    // ...
    await task_comletion_source.Task;
}

Обратите внимание, что это вызовет исключение в вашем ожидании, которое вы должны обработать.

14
Evk

Я рекомендую вам не строить это самостоятельно. Есть несколько случаев Edge вокруг токенов отмены, которые утомительны, чтобы получить право. Например, если регистрация, возвращенная из Register, никогда не удаляется, вы можете получить утечку ресурса.

Вместо этого вы можете использовать метод расширения Task.WaitAsync из моей библиотеки AsyncEx.Tasks :

var task_completion_source = new TaskCompletionSource<bool>();
observable.Subscribe(b => 
{ 
  if (b) 
    task_completion_source.SetResult(true); 
});
await task_completion_source.Task.WaitAsync(cancellationToken);

Кстати, я настоятельно рекомендую вам использовать ToTask вместо явного TaskCompletionSource. Опять же, ToTask прекрасно обрабатывает случаи Edge для вас.

6
Stephen Cleary

Это был мой шанс написать это сам. Я почти совершил ошибку за то, что не избавился от Регистра (спасибо Стивену Клири)

    /// <summary>
    /// This allows a TaskCompletionSource to be await with a cancellation token and timeout.
    /// 
    /// Example usable:
    /// 
    ///     var tcs = new TaskCompletionSource<bool>();
    ///           ...
    ///     var result = await tcs.WaitAsync(timeoutTokenSource.Token);
    /// 
    /// A TaskCanceledException will be thrown if the given cancelToken is canceled before the tcs completes or errors. 
    /// </summary>
    /// <typeparam name="TResult">Result type of the TaskCompletionSource</typeparam>
    /// <param name="tcs">The task completion source to be used  </param>
    /// <param name="cancelToken">This method will throw an OperationCanceledException if the cancelToken is canceled</param>
    /// <param name="timeoutMs">This method will throw a TimeoutException if it doesn't complete within the given timeout, unless the timeout is less then or equal to 0 or Timeout.Infinite</param>
    /// <param name="updateTcs">If this is true and the given cancelToken is canceled then the underlying tcs will also be canceled.  If this is true a timeout occurs the underlying tcs will be faulted with a TimeoutException.</param>
    /// <returns>The tcs.Task</returns>
    public static async Task<TResult> WaitAsync<TResult>(this TaskCompletionSource<TResult> tcs, CancellationToken cancelToken, int timeoutMs = Timeout.Infinite, bool updateTcs = false)
    {
        // The overrideTcs is used so we can wait for either the give tcs to complete or the overrideTcs.  We do this using the Task.WhenAny method.
        // one issue with WhenAny is that it won't return when a task is canceled, it only returns when a task completes so we complete the
        // overrideTcs when either the cancelToken is canceled or the timeoutMs is reached.
        //
        var overrideTcs = new TaskCompletionSource<TResult>();
        using( var timeoutCancelTokenSource = (timeoutMs <= 0 || timeoutMs == Timeout.Infinite) ? null : new CancellationTokenSource(timeoutMs) )
        {
            var timeoutToken = timeoutCancelTokenSource?.Token ?? CancellationToken.None;
            using( var linkedTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancelToken, timeoutToken) )
            {
                // This method is called when either the linkedTokenSource is canceled.  This lets us assign a value to the overrideTcs so that
                // We can break out of the await WhenAny below.
                //
                void CancelTcs()
                {
                    if( updateTcs && !tcs.Task.IsCompleted )
                    {
                        // ReSharper disable once AccessToDisposedClosure (in this case, CancelTcs will never be called outside the using)
                        if( timeoutCancelTokenSource?.IsCancellationRequested ?? false )
                            tcs.TrySetException(new TimeoutException($"WaitAsync timed out after {timeoutMs}ms"));
                        else
                            tcs.TrySetCanceled();
                    }

                    overrideTcs.TrySetResult(default(TResult));
                }

                using( linkedTokenSource.Token.Register(CancelTcs) )
                {
                    try
                    {
                        await Task.WhenAny(tcs.Task, overrideTcs.Task);
                    }
                    catch { /* ignore */ }

                    // We always favor the result from the given tcs task if it has completed.
                    //
                    if( tcs.Task.IsCompleted )
                    {
                        // We do another await here so that if the tcs.Task has faulted or has been canceled we won't wrap those exceptions
                        // in a nested exception.  While technically accessing the tcs.Task.Result will generate the same exception the
                        // exception will be wrapped in a nested exception.  We don't want that nesting so we just await.
                        await tcs.Task;
                        return tcs.Task.Result;
                    }

                    // It wasn't the tcs.Task that got us our of the above WhenAny so go ahead and timeout or cancel the operation.
                    //
                    if( timeoutCancelTokenSource?.IsCancellationRequested ?? false )
                        throw new TimeoutException($"WaitAsync timed out after {timeoutMs}ms");

                    throw new OperationCanceledException();
                }
            }
        }
    }

Это с броском TaskCanceledException, если cancelToken отменен прежде, чем tcs получит результат или ошибки.

1
Tod Cunningham