Improve cancellation (use ours).
MODIFY
WalletWasabi.Tests/IntegrationTests/P2pTests.cs
WalletWasabi.Tests/IntegrationTests/P2pTests.cs
@@ -35,7 +35,7 @@ namespace WalletWasabi.Tests.IntegrationTests
{
await RuntimeParams.LoadAsync();
var network = Network.GetNetwork(networkString);
- var blocksToDownload = new HashSet<uint256>();
+ var blocksToDownload = new List<uint256>();
if (network == Network.Main)
{
blocksToDownload.Add(new uint256("00000000000000000037c2de35bd85f3e57f14ddd741ce6cee5b28e51473d5d0"));
@@ -107,26 +107,41 @@ namespace WalletWasabi.Tests.IntegrationTests
var mempoolTransactionAwaiter = new EventAwaiter<SmartTransaction>(
h => bitcoinStore.MempoolService.TransactionReceived += h,
h => bitcoinStore.MempoolService.TransactionReceived -= h);
- using var mempoolTransactionReceivedTimeoutCts = new CancellationTokenSource(TimeSpan.FromMinutes(1));
var nodeConnectionAwaiter = new EventAwaiter<NodeEventArgs>(
h => nodes.ConnectedNodes.Added += h,
h => nodes.ConnectedNodes.Added -= h);
- using var nodeAddedTimeoutCts = new CancellationTokenSource(TimeSpan.FromMinutes(1));
nodes.Connect();
+ await nodeConnectionAwaiter.Task.WithAwaitCancellationAsync(TimeSpan.FromMinutes(1));
+
+ nodeConnectionAwaiter = new EventAwaiter<NodeEventArgs>(
+ h => nodes.ConnectedNodes.Added += h,
+ h => nodes.ConnectedNodes.Added -= h);
+ await nodeConnectionAwaiter.Task.WithAwaitCancellationAsync(TimeSpan.FromMinutes(1));
+
+ nodeConnectionAwaiter = new EventAwaiter<NodeEventArgs>(
+ h => nodes.ConnectedNodes.Added += h,
+ h => nodes.ConnectedNodes.Added -= h);
+ await nodeConnectionAwaiter.Task.WithAwaitCancellationAsync(TimeSpan.FromMinutes(1));
+
+ var downloadTasks = new List<Task<Block>>();
foreach (var hash in blocksToDownload)
{
using var cts = new CancellationTokenSource(TimeSpan.FromMinutes(3));
- var block = await walletService.FetchBlockAsync(hash, cts.Token);
- Assert.True(File.Exists(Path.Combine(blocksFolderPath, hash.ToString())));
- Logger.LogInfo($"Full block is downloaded: {hash}.");
+ downloadTasks.Add(walletService.FetchBlockAsync(hash, cts.Token));
+ }
+
+ var i = 0;
+ var hashArray = blocksToDownload.ToArray();
+ foreach (var block in await Task.WhenAll(downloadTasks))
+ {
+ Assert.True(File.Exists(Path.Combine(blocksFolderPath, hashArray[i].ToString())));
+ i++;
}
- await Task.WhenAll(
- mempoolTransactionAwaiter.Task.WithCancellation(mempoolTransactionReceivedTimeoutCts.Token),
- nodeConnectionAwaiter.Task.WithCancellation(nodeAddedTimeoutCts.Token));
+ await mempoolTransactionAwaiter.Task.WithAwaitCancellationAsync(TimeSpan.FromMinutes(1));
}
finally
{

MODIFY
WalletWasabi.Tests/UnitTests/BitcoinCore/P2pBasedTests.cs
WalletWasabi.Tests/UnitTests/BitcoinCore/P2pBasedTests.cs
@@ -45,8 +45,7 @@ namespace WalletWasabi.Tests.UnitTests.BitcoinCore
var txid = await rpc.SendToAddressAsync(addr, Money.Coins(1));
Assert.NotNull(txid);
- using var cts = new CancellationTokenSource(1000);
- var stx = await eventAwaiter.Task.WithCancellation(cts.Token);
+ var stx = await eventAwaiter.Task.WithAwaitCancellationAsync(1000);
Assert.Equal(txid, stx.GetHash());
}

MODIFY
WalletWasabi/Extensions/TaskExtensions.cs
WalletWasabi/Extensions/TaskExtensions.cs
@@ -6,7 +6,7 @@ namespace System.Threading.Tasks
/// UNSAFE! You are cancelling the wait on the callback of the original Task, not cancelling the operation itself.
/// https://stackoverflow.com/questions/14524209/what-is-the-correct-way-to-cancel-an-async-operation-that-doesnt-accept-a-cance/14524565#14524565
/// </summary>
- public static async Task<T> WithAwaitCancellationAsync<T>(this Task<T> me, CancellationToken cancel, int waitForGracefulTermination = 0)
+ public static async Task<T> WithAwaitCancellationAsync<T>(this Task<T> me, CancellationToken cancel, int waitForGracefulTerminationMilliseconds = 0)
{
// The tasck completion source.
var tcs = new TaskCompletionSource<bool>();
@@ -16,22 +16,67 @@ namespace System.Threading.Tasks
s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
{
// If the task waited on is the cancellation token...
- if (me != await Task.WhenAny(me, tcs.Task))
+ if (me != await Task.WhenAny(me, tcs.Task).ConfigureAwait(false))
{
- if (waitForGracefulTermination <= 0)
+ if (waitForGracefulTerminationMilliseconds <= 0)
{
throw new OperationCanceledException(cancel);
}
else
{
- using var cts = new CancellationTokenSource(waitForGracefulTermination);
- return await me.WithAwaitCancellationAsync(cts.Token);
+ using var cts = new CancellationTokenSource(waitForGracefulTerminationMilliseconds);
+ return await me.WithAwaitCancellationAsync(cts.Token).ConfigureAwait(false);
}
}
}
// Wait for one or the other to complete.
- return await me;
+ return await me.ConfigureAwait(false);
}
+
+ /// <summary>
+ /// UNSAFE! You are cancelling the wait on the callback of the original Task, not cancelling the operation itself.
+ /// https://stackoverflow.com/questions/14524209/what-is-the-correct-way-to-cancel-an-async-operation-that-doesnt-accept-a-cance/14524565#14524565
+ /// </summary>
+ public static async Task WithAwaitCancellationAsync(this Task me, CancellationToken cancel, int waitForGracefulTerminationMilliseconds = 0)
+ {
+ // The tasck completion source.
+ var tcs = new TaskCompletionSource<bool>();
+
+ // Register with the cancellation token.
+ using (cancel.Register(
+ s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
+ {
+ // If the task waited on is the cancellation token...
+ if (me != await Task.WhenAny(me, tcs.Task).ConfigureAwait(false))
+ {
+ if (waitForGracefulTerminationMilliseconds <= 0)
+ {
+ throw new OperationCanceledException(cancel);
+ }
+ else
+ {
+ using var cts = new CancellationTokenSource(waitForGracefulTerminationMilliseconds);
+ await me.WithAwaitCancellationAsync(cts.Token).ConfigureAwait(false);
+ return;
+ }
+ }
+ }
+
+ // Wait for one or the other to complete.
+ await me.ConfigureAwait(false);
+ }
+
+ public static async Task<T> WithAwaitCancellationAsync<T>(this Task<T> me, int millisecondsDelay)
+ => await me.WithAwaitCancellationAsync(CancellationToken.None, millisecondsDelay).ConfigureAwait(false);
+
+ public static async Task<T> WithAwaitCancellationAsync<T>(this Task<T> me, TimeSpan timeout)
+ => await me.WithAwaitCancellationAsync(CancellationToken.None, (int)timeout.TotalMilliseconds).ConfigureAwait(false);
+
+ public static async Task WithAwaitCancellationAsync(this Task me, int millisecondsDelay)
+ => await me.WithAwaitCancellationAsync(CancellationToken.None, millisecondsDelay).ConfigureAwait(false);
+
+ public static async Task WithAwaitCancellationAsync(this Task me, TimeSpan timeout)
+ => await me.WithAwaitCancellationAsync(CancellationToken.None, (int)timeout.TotalMilliseconds).ConfigureAwait(false);
}
}

MODIFY
WalletWasabi/Services/WalletService.cs
WalletWasabi/Services/WalletService.cs
@@ -865,10 +865,7 @@ namespace WalletWasabi.Services
}
var invPayload = new InvPayload(transaction.Transaction);
// Give 7 seconds to send the inv payload.
- using (var cts = new CancellationTokenSource(TimeSpan.FromSeconds(7)))
- {
- await node.SendMessageAsync(invPayload).WithCancellation(cts.Token); // ToDo: It's dangerous way to cancel. Implement proper cancellation to NBitcoin!
- }
+ await node.SendMessageAsync(invPayload).WithAwaitCancellationAsync(TimeSpan.FromSeconds(7)); // ToDo: It's dangerous way to cancel. Implement proper cancellation to NBitcoin!
if (BitcoinStore.MempoolService.TryGetFromBroadcastStore(transaction.GetHash(), out TransactionBroadcastEntry entry))
{

Improve cancellation (use ours).