Fixed some query cancellation implementation details (#124)

Fixed one bug where the wrong cancellation token was passed to the async query impl...
Removed the custom cancellation callback to dispose UdpClient for example, can now use the cancellation token + a callback registration.
Includes other fix from #121

Author: MichaCo

samples/MiniDig/RandomCommand.cs

@@ -73,7 +73,7 @@ protected override async Task<int> Execute()
             _runSync = SyncArg.HasValue();
 
             _settings = GetLookupSettings();
-            _settings.EnableAuditTrail = true;
+            _settings.EnableAuditTrail = false;
             _settings.ThrowDnsErrors = false;
             _settings.ContinueOnDnsError = false;
             _lookup = GetDnsLookup(_settings);

src/DnsClient/DnsMessageHandler.cs

@@ -23,7 +23,6 @@ internal abstract class DnsMessageHandler
         public abstract Task<DnsResponseMessage> QueryAsync(
             IPEndPoint endpoint,
             DnsRequestMessage request,
-            Action<Action> cancelationCallback,
             CancellationToken cancellationToken);
 
         // Transient errors will be retried on the same NameServer before the resolver moves on

src/DnsClient/DnsTcpMessageHandler.cs

@@ -23,20 +23,18 @@ public override DnsResponseMessage Query(IPEndPoint endpoint, DnsRequestMessage
             {
                 using (var cts = new CancellationTokenSource(timeout))
                 {
-                    Action onCancel = () => { };
-                    return QueryAsync(endpoint, request, (s) => onCancel = s, cts.Token)
-                        .WithCancellation(onCancel, cts.Token)
+                    return QueryAsync(endpoint, request, cts.Token)
+                        .WithCancellation(cts.Token)
                         .ConfigureAwait(false).GetAwaiter().GetResult();
                 }
             }
 
-            return QueryAsync(endpoint, request, (s) => { }, CancellationToken.None).ConfigureAwait(false).GetAwaiter().GetResult();
+            return QueryAsync(endpoint, request, CancellationToken.None).ConfigureAwait(false).GetAwaiter().GetResult();
         }
 
         public override async Task<DnsResponseMessage> QueryAsync(
             IPEndPoint server,
             DnsRequestMessage request,
-            Action<Action> cancelationCallback,
             CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
@@ -48,7 +46,7 @@ public override async Task<DnsResponseMessage> QueryAsync(
                 _pools.TryAdd(server, new ClientPool(true, server));
             }
 
-            cancelationCallback(() =>
+            using var cancelCallback = cancellationToken.Register(() =>
             {
                 if (entry == null)
                 {

src/DnsClient/DnsUdpMessageHandler.cs

@@ -82,26 +82,24 @@ public override DnsResponseMessage Query(
         public override async Task<DnsResponseMessage> QueryAsync(
             IPEndPoint endpoint,
             DnsRequestMessage request,
-            Action<Action> cancelationCallback,
             CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
 
             UdpClient udpClient = GetNextUdpClient(endpoint.AddressFamily);
 
-            bool mustDispose = false;
-            try
+            using var callback = cancellationToken.Register(() =>
             {
-                // setup timeout cancellation, dispose socket (the only way to actually cancel the request in async...
-                cancelationCallback(() =>
-                {
 #if !NET45
-                    udpClient.Dispose();
+                udpClient.Dispose();
 #else
-                    udpClient.Close();
+                udpClient.Close();
 #endif
-                });
+            });
 
+            bool mustDispose = false;
+            try
+            {
                 using (var writer = new DnsDatagramWriter())
                 {
                     GetRequestData(request, writer);
@@ -128,6 +126,10 @@ public override async Task<DnsResponseMessage> QueryAsync(
                     return response;
                 }
             }
+            catch (SocketException se) when (se.SocketErrorCode == SocketError.OperationAborted)
+            {
+                throw new TimeoutException();
+            }
             catch (ObjectDisposedException)
             {
                 // we disposed it in case of a timeout request, lets indicate it actually timed out...

src/DnsClient/LookupClient.cs

@@ -1097,15 +1097,6 @@ private async Task<IDnsQueryResponse> ResolveQueryAsync(
                         audit?.StartTimer();
 
                         DnsResponseMessage response;
-                        Action onCancel = () => { };
-                        Task<DnsResponseMessage> resultTask = handler.QueryAsync(
-                            serverInfo.IPEndPoint,
-                            request,
-                            (cancel) =>
-                            {
-                                onCancel = cancel;
-                            },
-                            cancellationToken);
 
                         if (settings.Timeout != System.Threading.Timeout.InfiniteTimeSpan
                             || (cancellationToken != CancellationToken.None && cancellationToken.CanBeCanceled))
@@ -1119,12 +1110,21 @@ private async Task<IDnsQueryResponse> ResolveQueryAsync(
                             using (cts)
                             using (linkedCts)
                             {
-                                response = await resultTask.WithCancellation(onCancel, (linkedCts ?? cts).Token).ConfigureAwait(false);
+                                response = await handler.QueryAsync(
+                                    serverInfo.IPEndPoint,
+                                    request,
+                                    (linkedCts ?? cts).Token)
+                                .WithCancellation((linkedCts ?? cts).Token)
+                                .ConfigureAwait(false);
                             }
                         }
                         else
                         {
-                            response = await resultTask.ConfigureAwait(false);
+                            response = await handler.QueryAsync(
+                                    serverInfo.IPEndPoint,
+                                    request,
+                                    cancellationToken)
+                                .ConfigureAwait(false);
                         }
 
                         lastQueryResponse = ProcessResponseMessage(

src/DnsClient/TaskExtensions.cs

@@ -2,19 +2,17 @@
 {
     internal static class TaskExtensions
     {
-        public static async Task<T> WithCancellation<T>(this Task<T> task, Action onCancel, CancellationToken cancellationToken)
+        public static async Task<T> WithCancellation<T>(this Task<T> task, CancellationToken cancellationToken)
         {
             var tcs = new TaskCompletionSource<bool>();
 
             using (cancellationToken.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
             {
                 if (task != await Task.WhenAny(task, tcs.Task).ConfigureAwait(false))
                 {
-                    try
-                    {
-                        onCancel();
-                    }
-                    catch { }
+                    // observe the exception to avoid "System.AggregateException: A Task's exception(s) were
+                    // not observed either by Waiting on the Task or accessing its Exception property."
+                    _ = task.ContinueWith(t => t.Exception);
                     throw new OperationCanceledException(cancellationToken);
                 }
             }

test-other/Benchmarks/DnsClientBenchmarks.DatagramReader.cs

@@ -138,7 +138,6 @@ public override DnsResponseMessage Query(
                 public override Task<DnsResponseMessage> QueryAsync(
                     IPEndPoint server,
                     DnsRequestMessage request,
-                    Action<Action> cancelationCallback,
                     CancellationToken cancellationToken)
                 {
                     // no need to run async here as we don't do any IO

test-other/Benchmarks/DnsClientBenchmarks.RequestResponseParsing.cs

@@ -166,7 +166,6 @@ public override DnsResponseMessage Query(
         public override Task<DnsResponseMessage> QueryAsync(
             IPEndPoint server,
             DnsRequestMessage request,
-            Action<Action> cancelationCallback,
             CancellationToken cancellationToken)
         {
             return Task.FromResult(Query(server, request, Timeout.InfiniteTimeSpan));

test/DnsClient.Tests/DnsClient.Tests.csproj

@@ -9,6 +9,7 @@
     <PublicSign Condition=" '$(OS)' != 'Windows_NT' ">true</PublicSign>
     <PackageId>DnsClient.Tests</PackageId>
     <IsPackable>false</IsPackable>
+    <LangVersion>latest</LangVersion>
   </PropertyGroup>
 
   <PropertyGroup>

test/DnsClient.Tests/DnsResponseParsingTest.cs

@@ -359,7 +359,6 @@ public override DnsResponseMessage Query(
             public override Task<DnsResponseMessage> QueryAsync(
                 IPEndPoint server,
                 DnsRequestMessage request,
-                Action<Action> cancelationCallback,
                 CancellationToken cancellationToken)
             {
                 // no need to run async here as we don't do any IO