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
MichaCo · Jun 12, 2021 · de7fc4c · de7fc4c
1 parent 314f0a7
commit de7fc4c
Show file tree

Hide file tree

Showing 12 changed files with 33 additions and 39 deletions.
diff --git a/samples/MiniDig/RandomCommand.cs b/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);

diff --git a/src/DnsClient/DnsMessageHandler.cs b/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

diff --git a/src/DnsClient/DnsTcpMessageHandler.cs b/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)
                 {

diff --git a/src/DnsClient/DnsUdpMessageHandler.cs b/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...

diff --git a/src/DnsClient/LookupClient.cs b/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(

diff --git a/src/DnsClient/TaskExtensions.cs b/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);
                 }
             }

diff --git a/test-other/Benchmarks/DnsClientBenchmarks.DatagramReader.cs b/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

diff --git a/test-other/Benchmarks/DnsClientBenchmarks.RequestResponseParsing.cs b/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));

diff --git a/test/DnsClient.Tests/DnsClient.Tests.csproj b/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>

diff --git a/test/DnsClient.Tests/DnsResponseParsingTest.cs b/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

diff --git a/test/DnsClient.Tests/LookupClientRetryTest.cs b/test/DnsClient.Tests/LookupClientRetryTest.cs
@@ -1664,7 +1664,7 @@ public override DnsResponseMessage Query(IPEndPoint endpoint, DnsRequestMessage
             return _onQuery(endpoint, request);
         }
 
-        public override Task<DnsResponseMessage> QueryAsync(IPEndPoint endpoint, DnsRequestMessage request, Action<Action> cancelationCallback, CancellationToken cancellationToken)
+        public override Task<DnsResponseMessage> QueryAsync(IPEndPoint endpoint, DnsRequestMessage request, CancellationToken cancellationToken)
         {
             return Task.FromResult(_onQuery(endpoint, request));
         }

diff --git a/test/DnsClient.Tests/LookupConfigurationTest.cs b/test/DnsClient.Tests/LookupConfigurationTest.cs
@@ -1039,7 +1039,6 @@ public override DnsResponseMessage Query(
             public override Task<DnsResponseMessage> QueryAsync(
                 IPEndPoint server,
                 DnsRequestMessage request,
-                Action<Action> cancelationCallback,
                 CancellationToken cancellationToken)
             {
                 LastServer = server;