Skip to content

Commit de7fc4c

Browse files
authored
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
1 parent 314f0a7 commit de7fc4c

12 files changed

+33
-39
lines changed

samples/MiniDig/RandomCommand.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ protected override async Task<int> Execute()
7373
_runSync = SyncArg.HasValue();
7474

7575
_settings = GetLookupSettings();
76-
_settings.EnableAuditTrail = true;
76+
_settings.EnableAuditTrail = false;
7777
_settings.ThrowDnsErrors = false;
7878
_settings.ContinueOnDnsError = false;
7979
_lookup = GetDnsLookup(_settings);

src/DnsClient/DnsMessageHandler.cs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@ internal abstract class DnsMessageHandler
2323
public abstract Task<DnsResponseMessage> QueryAsync(
2424
IPEndPoint endpoint,
2525
DnsRequestMessage request,
26-
Action<Action> cancelationCallback,
2726
CancellationToken cancellationToken);
2827

2928
// Transient errors will be retried on the same NameServer before the resolver moves on

src/DnsClient/DnsTcpMessageHandler.cs

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,20 +23,18 @@ public override DnsResponseMessage Query(IPEndPoint endpoint, DnsRequestMessage
2323
{
2424
using (var cts = new CancellationTokenSource(timeout))
2525
{
26-
Action onCancel = () => { };
27-
return QueryAsync(endpoint, request, (s) => onCancel = s, cts.Token)
28-
.WithCancellation(onCancel, cts.Token)
26+
return QueryAsync(endpoint, request, cts.Token)
27+
.WithCancellation(cts.Token)
2928
.ConfigureAwait(false).GetAwaiter().GetResult();
3029
}
3130
}
3231

33-
return QueryAsync(endpoint, request, (s) => { }, CancellationToken.None).ConfigureAwait(false).GetAwaiter().GetResult();
32+
return QueryAsync(endpoint, request, CancellationToken.None).ConfigureAwait(false).GetAwaiter().GetResult();
3433
}
3534

3635
public override async Task<DnsResponseMessage> QueryAsync(
3736
IPEndPoint server,
3837
DnsRequestMessage request,
39-
Action<Action> cancelationCallback,
4038
CancellationToken cancellationToken)
4139
{
4240
cancellationToken.ThrowIfCancellationRequested();
@@ -48,7 +46,7 @@ public override async Task<DnsResponseMessage> QueryAsync(
4846
_pools.TryAdd(server, new ClientPool(true, server));
4947
}
5048

51-
cancelationCallback(() =>
49+
using var cancelCallback = cancellationToken.Register(() =>
5250
{
5351
if (entry == null)
5452
{

src/DnsClient/DnsUdpMessageHandler.cs

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -82,26 +82,24 @@ public override DnsResponseMessage Query(
8282
public override async Task<DnsResponseMessage> QueryAsync(
8383
IPEndPoint endpoint,
8484
DnsRequestMessage request,
85-
Action<Action> cancelationCallback,
8685
CancellationToken cancellationToken)
8786
{
8887
cancellationToken.ThrowIfCancellationRequested();
8988

9089
UdpClient udpClient = GetNextUdpClient(endpoint.AddressFamily);
9190

92-
bool mustDispose = false;
93-
try
91+
using var callback = cancellationToken.Register(() =>
9492
{
95-
// setup timeout cancellation, dispose socket (the only way to actually cancel the request in async...
96-
cancelationCallback(() =>
97-
{
9893
#if !NET45
99-
udpClient.Dispose();
94+
udpClient.Dispose();
10095
#else
101-
udpClient.Close();
96+
udpClient.Close();
10297
#endif
103-
});
98+
});
10499

100+
bool mustDispose = false;
101+
try
102+
{
105103
using (var writer = new DnsDatagramWriter())
106104
{
107105
GetRequestData(request, writer);
@@ -128,6 +126,10 @@ public override async Task<DnsResponseMessage> QueryAsync(
128126
return response;
129127
}
130128
}
129+
catch (SocketException se) when (se.SocketErrorCode == SocketError.OperationAborted)
130+
{
131+
throw new TimeoutException();
132+
}
131133
catch (ObjectDisposedException)
132134
{
133135
// we disposed it in case of a timeout request, lets indicate it actually timed out...

src/DnsClient/LookupClient.cs

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1097,15 +1097,6 @@ private async Task<IDnsQueryResponse> ResolveQueryAsync(
10971097
audit?.StartTimer();
10981098

10991099
DnsResponseMessage response;
1100-
Action onCancel = () => { };
1101-
Task<DnsResponseMessage> resultTask = handler.QueryAsync(
1102-
serverInfo.IPEndPoint,
1103-
request,
1104-
(cancel) =>
1105-
{
1106-
onCancel = cancel;
1107-
},
1108-
cancellationToken);
11091100

11101101
if (settings.Timeout != System.Threading.Timeout.InfiniteTimeSpan
11111102
|| (cancellationToken != CancellationToken.None && cancellationToken.CanBeCanceled))
@@ -1119,12 +1110,21 @@ private async Task<IDnsQueryResponse> ResolveQueryAsync(
11191110
using (cts)
11201111
using (linkedCts)
11211112
{
1122-
response = await resultTask.WithCancellation(onCancel, (linkedCts ?? cts).Token).ConfigureAwait(false);
1113+
response = await handler.QueryAsync(
1114+
serverInfo.IPEndPoint,
1115+
request,
1116+
(linkedCts ?? cts).Token)
1117+
.WithCancellation((linkedCts ?? cts).Token)
1118+
.ConfigureAwait(false);
11231119
}
11241120
}
11251121
else
11261122
{
1127-
response = await resultTask.ConfigureAwait(false);
1123+
response = await handler.QueryAsync(
1124+
serverInfo.IPEndPoint,
1125+
request,
1126+
cancellationToken)
1127+
.ConfigureAwait(false);
11281128
}
11291129

11301130
lastQueryResponse = ProcessResponseMessage(

src/DnsClient/TaskExtensions.cs

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,17 @@
22
{
33
internal static class TaskExtensions
44
{
5-
public static async Task<T> WithCancellation<T>(this Task<T> task, Action onCancel, CancellationToken cancellationToken)
5+
public static async Task<T> WithCancellation<T>(this Task<T> task, CancellationToken cancellationToken)
66
{
77
var tcs = new TaskCompletionSource<bool>();
88

99
using (cancellationToken.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
1010
{
1111
if (task != await Task.WhenAny(task, tcs.Task).ConfigureAwait(false))
1212
{
13-
try
14-
{
15-
onCancel();
16-
}
17-
catch { }
13+
// observe the exception to avoid "System.AggregateException: A Task's exception(s) were
14+
// not observed either by Waiting on the Task or accessing its Exception property."
15+
_ = task.ContinueWith(t => t.Exception);
1816
throw new OperationCanceledException(cancellationToken);
1917
}
2018
}

test-other/Benchmarks/DnsClientBenchmarks.DatagramReader.cs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,6 @@ public override DnsResponseMessage Query(
138138
public override Task<DnsResponseMessage> QueryAsync(
139139
IPEndPoint server,
140140
DnsRequestMessage request,
141-
Action<Action> cancelationCallback,
142141
CancellationToken cancellationToken)
143142
{
144143
// no need to run async here as we don't do any IO

test-other/Benchmarks/DnsClientBenchmarks.RequestResponseParsing.cs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,6 @@ public override DnsResponseMessage Query(
166166
public override Task<DnsResponseMessage> QueryAsync(
167167
IPEndPoint server,
168168
DnsRequestMessage request,
169-
Action<Action> cancelationCallback,
170169
CancellationToken cancellationToken)
171170
{
172171
return Task.FromResult(Query(server, request, Timeout.InfiniteTimeSpan));

test/DnsClient.Tests/DnsClient.Tests.csproj

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
<PublicSign Condition=" '$(OS)' != 'Windows_NT' ">true</PublicSign>
1010
<PackageId>DnsClient.Tests</PackageId>
1111
<IsPackable>false</IsPackable>
12+
<LangVersion>latest</LangVersion>
1213
</PropertyGroup>
1314

1415
<PropertyGroup>

test/DnsClient.Tests/DnsResponseParsingTest.cs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -359,7 +359,6 @@ public override DnsResponseMessage Query(
359359
public override Task<DnsResponseMessage> QueryAsync(
360360
IPEndPoint server,
361361
DnsRequestMessage request,
362-
Action<Action> cancelationCallback,
363362
CancellationToken cancellationToken)
364363
{
365364
// no need to run async here as we don't do any IO

0 commit comments

Comments
 (0)