From e808da9df92f52df4198bd0f56ce5e818f27b750 Mon Sep 17 00:00:00 2001 From: zs311521 <153907395+zs311521@users.noreply.github.com> Date: Sun, 9 Aug 2026 22:01:55 +1000 Subject: [PATCH] Fix bounded DNS-over-QUIC stream lifetime --- .../ClientConnection/DnsClientConnection.cs | 592 +++++++++++-- .../ClientConnection/QuicClientConnection.cs | 807 ++++++++++++++++-- 2 files changed, 1245 insertions(+), 154 deletions(-) diff --git a/TechnitiumLibrary.Net/Dns/ClientConnection/DnsClientConnection.cs b/TechnitiumLibrary.Net/Dns/ClientConnection/DnsClientConnection.cs index 27508e70..f7c61969 100644 --- a/TechnitiumLibrary.Net/Dns/ClientConnection/DnsClientConnection.cs +++ b/TechnitiumLibrary.Net/Dns/ClientConnection/DnsClientConnection.cs @@ -57,6 +57,11 @@ public abstract class DnsClientConnection : IDisposable, IAsyncDisposable static readonly ConcurrentDictionary> _existingHttpsConnections = new ConcurrentDictionary>(); static readonly ConcurrentDictionary> _existingQuicConnections = new ConcurrentDictionary>(); + static readonly ReaderWriterLockSlim _tcpConnectionsLock = new ReaderWriterLockSlim(LockRecursionPolicy.NoRecursion); + static readonly ReaderWriterLockSlim _tlsConnectionsLock = new ReaderWriterLockSlim(LockRecursionPolicy.NoRecursion); + static readonly ReaderWriterLockSlim _httpsConnectionsLock = new ReaderWriterLockSlim(LockRecursionPolicy.NoRecursion); + static readonly ReaderWriterLockSlim _quicConnectionsLock = new ReaderWriterLockSlim(LockRecursionPolicy.NoRecursion); + #endregion #region constructor @@ -74,18 +79,45 @@ static DnsClientConnection() { foreach (KeyValuePair connection in existingTcpConnection.Value) { - if (connection.Value.LastQueried < expiryTime) + TcpClientConnection removedConnection = null; + + _tcpConnectionsLock.EnterWriteLock(); + try { - if (existingTcpConnection.Value.TryRemove(connection.Key, out TcpClientConnection removedConnection)) + if (connection.Value.CanEvictPooledConnection() && (connection.Value.LastQueried < expiryTime)) { - removedConnection.Pooled = false; - removedConnection.Dispose(); + if (!existingTcpConnection.Value.TryRemove(connection.Key, out removedConnection) || !ReferenceEquals(removedConnection, connection.Value) || !removedConnection.TryBeginPooledEviction()) + throw new InvalidOperationException("Failed to reserve the exact pooled TCP connection for eviction."); + } + } + finally + { + _tcpConnectionsLock.ExitWriteLock(); + } + + if (removedConnection is not null) + { + try + { + await removedConnection.DisposePooledAsync(delegate { removedConnection.Pooled = false; }); + } + catch (Exception ex) + { + ReportPoolMaintenanceError(ex); } } } - if (existingTcpConnection.Value.IsEmpty) - _existingTcpConnections.TryRemove(existingTcpConnection.Key, out _); + _tcpConnectionsLock.EnterWriteLock(); + try + { + if (existingTcpConnection.Value.IsEmpty && _existingTcpConnections.TryGetValue(existingTcpConnection.Key, out ConcurrentDictionary currentTcpConnections) && ReferenceEquals(currentTcpConnections, existingTcpConnection.Value)) + _existingTcpConnections.TryRemove(existingTcpConnection.Key, out _); + } + finally + { + _tcpConnectionsLock.ExitWriteLock(); + } } //cleanup unused tls connections @@ -93,18 +125,45 @@ static DnsClientConnection() { foreach (KeyValuePair connection in existingTlsConnection.Value) { - if (connection.Value.LastQueried < expiryTime) + TlsClientConnection removedConnection = null; + + _tlsConnectionsLock.EnterWriteLock(); + try + { + if (connection.Value.CanEvictPooledConnection() && (connection.Value.LastQueried < expiryTime)) + { + if (!existingTlsConnection.Value.TryRemove(connection.Key, out removedConnection) || !ReferenceEquals(removedConnection, connection.Value) || !removedConnection.TryBeginPooledEviction()) + throw new InvalidOperationException("Failed to reserve the exact pooled TLS connection for eviction."); + } + } + finally + { + _tlsConnectionsLock.ExitWriteLock(); + } + + if (removedConnection is not null) { - if (existingTlsConnection.Value.TryRemove(connection.Key, out TlsClientConnection removedConnection)) + try + { + await removedConnection.DisposePooledAsync(delegate { removedConnection.Pooled = false; }); + } + catch (Exception ex) { - removedConnection.Pooled = false; - removedConnection.Dispose(); + ReportPoolMaintenanceError(ex); } } } - if (existingTlsConnection.Value.IsEmpty) - _existingTlsConnections.TryRemove(existingTlsConnection.Key, out _); + _tlsConnectionsLock.EnterWriteLock(); + try + { + if (existingTlsConnection.Value.IsEmpty && _existingTlsConnections.TryGetValue(existingTlsConnection.Key, out ConcurrentDictionary currentTlsConnections) && ReferenceEquals(currentTlsConnections, existingTlsConnection.Value)) + _existingTlsConnections.TryRemove(existingTlsConnection.Key, out _); + } + finally + { + _tlsConnectionsLock.ExitWriteLock(); + } } //cleanup unused https connections @@ -112,18 +171,45 @@ static DnsClientConnection() { foreach (KeyValuePair connection in existingHttpsConnection.Value) { - if (connection.Value.LastQueried < expiryTime) + HttpsClientConnection removedConnection = null; + + _httpsConnectionsLock.EnterWriteLock(); + try + { + if (connection.Value.CanEvictPooledConnection() && (connection.Value.LastQueried < expiryTime)) + { + if (!existingHttpsConnection.Value.TryRemove(connection.Key, out removedConnection) || !ReferenceEquals(removedConnection, connection.Value) || !removedConnection.TryBeginPooledEviction()) + throw new InvalidOperationException("Failed to reserve the exact pooled HTTPS connection for eviction."); + } + } + finally + { + _httpsConnectionsLock.ExitWriteLock(); + } + + if (removedConnection is not null) { - if (existingHttpsConnection.Value.TryRemove(connection.Key, out HttpsClientConnection removedConnection)) + try + { + await removedConnection.DisposePooledAsync(delegate { removedConnection.Pooled = false; }); + } + catch (Exception ex) { - removedConnection.Pooled = false; - removedConnection.Dispose(); + ReportPoolMaintenanceError(ex); } } } - if (existingHttpsConnection.Value.IsEmpty) - _existingHttpsConnections.TryRemove(existingHttpsConnection.Key, out _); + _httpsConnectionsLock.EnterWriteLock(); + try + { + if (existingHttpsConnection.Value.IsEmpty && _existingHttpsConnections.TryGetValue(existingHttpsConnection.Key, out ConcurrentDictionary currentHttpsConnections) && ReferenceEquals(currentHttpsConnections, existingHttpsConnection.Value)) + _existingHttpsConnections.TryRemove(existingHttpsConnection.Key, out _); + } + finally + { + _httpsConnectionsLock.ExitWriteLock(); + } } //cleanup unused quic connections @@ -131,22 +217,51 @@ static DnsClientConnection() { foreach (KeyValuePair connection in existingQuicConnection.Value) { - if (connection.Value.LastQueried < expiryTime) + QuicClientConnection removedConnection = null; + + _quicConnectionsLock.EnterWriteLock(); + try + { + if (connection.Value.CanEvictPooledConnection() && (connection.Value.LastQueried < expiryTime)) + { + if (!existingQuicConnection.Value.TryRemove(connection.Key, out removedConnection) || !ReferenceEquals(removedConnection, connection.Value) || !removedConnection.TryBeginPooledEviction()) + throw new InvalidOperationException("Failed to reserve the exact pooled QUIC connection for eviction."); + } + } + finally + { + _quicConnectionsLock.ExitWriteLock(); + } + + if (removedConnection is not null) { - if (existingQuicConnection.Value.TryRemove(connection.Key, out QuicClientConnection removedConnection)) + try + { + await removedConnection.DisposePooledAsync(delegate { removedConnection.Pooled = false; }); + } + catch (Exception ex) { - removedConnection.Pooled = false; - await removedConnection.DisposeAsync(); + ReportPoolMaintenanceError(ex); } } } - if (existingQuicConnection.Value.IsEmpty) - _existingQuicConnections.TryRemove(existingQuicConnection.Key, out _); + _quicConnectionsLock.EnterWriteLock(); + try + { + if (existingQuicConnection.Value.IsEmpty && _existingQuicConnections.TryGetValue(existingQuicConnection.Key, out ConcurrentDictionary currentQuicConnections) && ReferenceEquals(currentQuicConnections, existingQuicConnection.Value)) + _existingQuicConnections.TryRemove(existingQuicConnection.Key, out _); + } + finally + { + _quicConnectionsLock.ExitWriteLock(); + } } } - catch - { } + catch (Exception ex) + { + ReportPoolMaintenanceError(ex); + } }); _maintenanceTimer.Change(MAINTENANCE_TIMER_INITIAL_INTERVAL, MAINTENANCE_TIMER_PERIODIC_INTERVAL); @@ -162,7 +277,21 @@ protected DnsClientConnection(NameServerAddress server, NetProxy proxy) #region IDisposable - bool _disposed; + enum DisposeOperation + { + ReleasePooledLease, + OwnPhysicalDispose, + JoinPhysicalDispose + } + + readonly object _disposeLock = new object(); + + int _pooledLeaseCount; + bool _isPooledConnection; + bool _poolEvictionStarted; + bool _physicalDisposeStarted; + TaskCompletionSource _physicalDisposeCompleted; + Task _physicalDisposeTask; protected virtual void Dispose(bool disposing) { } @@ -172,27 +301,210 @@ protected virtual ValueTask DisposeAsyncCore() return ValueTask.CompletedTask; } + private bool TryAcquirePooledLease() + { + if (Volatile.Read(ref _poolEvictionStarted) || (Volatile.Read(ref _physicalDisposeTask) is not null)) + return false; + + _isPooledConnection = true; + + int leaseCount = Interlocked.Increment(ref _pooledLeaseCount); + if (leaseCount <= 0) + { + Interlocked.Decrement(ref _pooledLeaseCount); + throw new InvalidOperationException("The pooled DNS connection lease count overflowed."); + } + + return true; + } + + private bool CanEvictPooledConnection() + { + return Volatile.Read(ref _isPooledConnection) && (Volatile.Read(ref _pooledLeaseCount) == 0) && !Volatile.Read(ref _poolEvictionStarted) && (Volatile.Read(ref _physicalDisposeTask) is null); + } + + private bool TryBeginPooledEviction() + { + lock (_disposeLock) + { + if (!_isPooledConnection || (_pooledLeaseCount != 0) || _poolEvictionStarted || (_physicalDisposeTask is not null)) + return false; + + _poolEvictionStarted = true; + _physicalDisposeCompleted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + _physicalDisposeTask = _physicalDisposeCompleted.Task; + return true; + } + } + + private DisposeOperation BeginDispose(out Task physicalDisposeTask, out TaskCompletionSource physicalDisposeCompleted) + { + if (Volatile.Read(ref _isPooledConnection) && !Volatile.Read(ref _poolEvictionStarted)) + { + int leaseCount = Interlocked.Decrement(ref _pooledLeaseCount); + if (leaseCount < 0) + { + Interlocked.Increment(ref _pooledLeaseCount); + throw new InvalidOperationException("The pooled DNS connection lease count is already zero."); + } + + physicalDisposeTask = null; + physicalDisposeCompleted = null; + return DisposeOperation.ReleasePooledLease; + } + + lock (_disposeLock) + { + if (_physicalDisposeTask is not null) + { + physicalDisposeTask = _physicalDisposeTask; + physicalDisposeCompleted = null; + return DisposeOperation.JoinPhysicalDispose; + } + + _physicalDisposeStarted = true; + _physicalDisposeCompleted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + _physicalDisposeTask = _physicalDisposeCompleted.Task; + + physicalDisposeTask = null; + physicalDisposeCompleted = _physicalDisposeCompleted; + return DisposeOperation.OwnPhysicalDispose; + } + } + + private ValueTask DisposePooledAsync(Action releaseFromPool) + { + Task physicalDisposeTask; + TaskCompletionSource physicalDisposeCompleted; + + lock (_disposeLock) + { + if (!_poolEvictionStarted || (_physicalDisposeTask is null) || (_physicalDisposeCompleted is null)) + throw new InvalidOperationException("The pooled DNS connection was not reserved for eviction."); + + physicalDisposeTask = _physicalDisposeTask; + + if (_physicalDisposeStarted) + return new ValueTask(physicalDisposeTask); + + _physicalDisposeStarted = true; + physicalDisposeCompleted = _physicalDisposeCompleted; + } + + _ = Task.Run(delegate { return CompletePooledDisposeAsync(releaseFromPool, physicalDisposeCompleted); }); + return new ValueTask(physicalDisposeTask); + } + + private async Task CompletePooledDisposeAsync(Action releaseFromPool, TaskCompletionSource physicalDisposeCompleted) + { + Exception disposeException = null; + + try + { + releaseFromPool(); + await DisposeAsyncCore().ConfigureAwait(false); + Dispose(false); + GC.SuppressFinalize(this); + } + catch (Exception ex) + { + disposeException = ex; + } + finally + { + CompletePhysicalDispose(physicalDisposeCompleted, disposeException); + } + } + + private async Task CompleteDirectDisposeAsync(TaskCompletionSource physicalDisposeCompleted) + { + Exception disposeException = null; + + try + { + await DisposeAsyncCore().ConfigureAwait(false); + Dispose(false); + GC.SuppressFinalize(this); + } + catch (Exception ex) + { + disposeException = ex; + } + finally + { + CompletePhysicalDispose(physicalDisposeCompleted, disposeException); + } + } + + private static void CompletePhysicalDispose(TaskCompletionSource physicalDisposeCompleted, Exception disposeException) + { + if (disposeException is null) + { + physicalDisposeCompleted.TrySetResult(); + } + else + { + physicalDisposeCompleted.TrySetException(disposeException); + _ = physicalDisposeCompleted.Task.Exception; + } + } + + private static void ReportPoolMaintenanceError(Exception exception) + { + try + { + Console.Error.WriteLine("DNS connection pool maintenance failed: " + exception); + } + catch + { } + } + public void Dispose() { - if (_disposed) + DisposeOperation operation = BeginDispose(out Task physicalDisposeTask, out TaskCompletionSource physicalDisposeCompleted); + + if (operation == DisposeOperation.ReleasePooledLease) + return; + + if (operation == DisposeOperation.JoinPhysicalDispose) + { + physicalDisposeTask.GetAwaiter().GetResult(); return; + } - Dispose(true); - GC.SuppressFinalize(this); + Exception disposeException = null; - _disposed = true; + try + { + Dispose(true); + GC.SuppressFinalize(this); + } + catch (Exception ex) + { + disposeException = ex; + throw; + } + finally + { + CompletePhysicalDispose(physicalDisposeCompleted, disposeException); + } } public async ValueTask DisposeAsync() { - if (_disposed) + DisposeOperation operation = BeginDispose(out Task physicalDisposeTask, out TaskCompletionSource physicalDisposeCompleted); + + if (operation == DisposeOperation.ReleasePooledLease) return; - await DisposeAsyncCore(); - Dispose(false); - GC.SuppressFinalize(this); + if (operation == DisposeOperation.JoinPhysicalDispose) + { + await physicalDisposeTask.ConfigureAwait(false); + return; + } - _disposed = true; + _ = Task.Run(delegate { return CompleteDirectDisposeAsync(physicalDisposeCompleted); }); + await physicalDisposeCompleted.Task.ConfigureAwait(false); } #endregion @@ -208,82 +520,202 @@ public static DnsClientConnection GetConnection(NameServerAddress server, NetPro case DnsTransportProtocol.Tcp: { - ConcurrentDictionary existingTcpConnection = _existingTcpConnections.GetOrAdd(server, delegate (NameServerAddress nameServer) - { - return new ConcurrentDictionary(); - }); + NetProxy proxyKey = proxy ?? NetProxy.NONE; - NetProxy proxyKey = proxy; + _tcpConnectionsLock.EnterReadLock(); + try + { + if (_existingTcpConnections.TryGetValue(server, out ConcurrentDictionary existingTcpConnection) && existingTcpConnection.TryGetValue(proxyKey, out TcpClientConnection connection)) + { + if (!connection.TryAcquirePooledLease()) + throw new InvalidOperationException("Failed to acquire the canonical pooled TCP connection."); - if (proxyKey is null) - proxyKey = NetProxy.NONE; + return connection; + } + } + finally + { + _tcpConnectionsLock.ExitReadLock(); + } - return existingTcpConnection.GetOrAdd(proxyKey, delegate (NetProxy netProxyKey) + _tcpConnectionsLock.EnterWriteLock(); + try { - TcpClientConnection connection = new TcpClientConnection(server, proxy); - connection.Pooled = true; + if (!_existingTcpConnections.TryGetValue(server, out ConcurrentDictionary existingTcpConnection)) + { + existingTcpConnection = new ConcurrentDictionary(); + if (!_existingTcpConnections.TryAdd(server, existingTcpConnection)) + throw new InvalidOperationException("Failed to add the canonical TCP connection pool."); + } + + if (!existingTcpConnection.TryGetValue(proxyKey, out TcpClientConnection connection)) + { + connection = new TcpClientConnection(server, proxy); + connection.Pooled = true; + + if (!existingTcpConnection.TryAdd(proxyKey, connection)) + throw new InvalidOperationException("Failed to add the canonical pooled TCP connection."); + } + + if (!connection.TryAcquirePooledLease()) + throw new InvalidOperationException("Failed to acquire the canonical pooled TCP connection."); + return connection; - }); + } + finally + { + _tcpConnectionsLock.ExitWriteLock(); + } } case DnsTransportProtocol.Tls: { - ConcurrentDictionary existingTlsConnection = _existingTlsConnections.GetOrAdd(server, delegate (NameServerAddress nameServer) - { - return new ConcurrentDictionary(); - }); + NetProxy proxyKey = proxy ?? NetProxy.NONE; - NetProxy proxyKey = proxy; + _tlsConnectionsLock.EnterReadLock(); + try + { + if (_existingTlsConnections.TryGetValue(server, out ConcurrentDictionary existingTlsConnection) && existingTlsConnection.TryGetValue(proxyKey, out TlsClientConnection connection)) + { + if (!connection.TryAcquirePooledLease()) + throw new InvalidOperationException("Failed to acquire the canonical pooled TLS connection."); - if (proxyKey is null) - proxyKey = NetProxy.NONE; + return connection; + } + } + finally + { + _tlsConnectionsLock.ExitReadLock(); + } - return existingTlsConnection.GetOrAdd(proxyKey, delegate (NetProxy netProxyKey) + _tlsConnectionsLock.EnterWriteLock(); + try { - TlsClientConnection connection = new TlsClientConnection(server, proxy); - connection.Pooled = true; + if (!_existingTlsConnections.TryGetValue(server, out ConcurrentDictionary existingTlsConnection)) + { + existingTlsConnection = new ConcurrentDictionary(); + if (!_existingTlsConnections.TryAdd(server, existingTlsConnection)) + throw new InvalidOperationException("Failed to add the canonical TLS connection pool."); + } + + if (!existingTlsConnection.TryGetValue(proxyKey, out TlsClientConnection connection)) + { + connection = new TlsClientConnection(server, proxy); + connection.Pooled = true; + + if (!existingTlsConnection.TryAdd(proxyKey, connection)) + throw new InvalidOperationException("Failed to add the canonical pooled TLS connection."); + } + + if (!connection.TryAcquirePooledLease()) + throw new InvalidOperationException("Failed to acquire the canonical pooled TLS connection."); + return connection; - }); + } + finally + { + _tlsConnectionsLock.ExitWriteLock(); + } } case DnsTransportProtocol.Https: { - ConcurrentDictionary existingHttpsConnection = _existingHttpsConnections.GetOrAdd(server, delegate (NameServerAddress nameServer) - { - return new ConcurrentDictionary(); - }); + NetProxy proxyKey = proxy ?? NetProxy.NONE; - NetProxy proxyKey = proxy; + _httpsConnectionsLock.EnterReadLock(); + try + { + if (_existingHttpsConnections.TryGetValue(server, out ConcurrentDictionary existingHttpsConnection) && existingHttpsConnection.TryGetValue(proxyKey, out HttpsClientConnection connection)) + { + if (!connection.TryAcquirePooledLease()) + throw new InvalidOperationException("Failed to acquire the canonical pooled HTTPS connection."); - if (proxyKey is null) - proxyKey = NetProxy.NONE; + return connection; + } + } + finally + { + _httpsConnectionsLock.ExitReadLock(); + } - return existingHttpsConnection.GetOrAdd(proxyKey, delegate (NetProxy netProxyKey) + _httpsConnectionsLock.EnterWriteLock(); + try { - HttpsClientConnection connection = new HttpsClientConnection(server, proxy); - connection.Pooled = true; + if (!_existingHttpsConnections.TryGetValue(server, out ConcurrentDictionary existingHttpsConnection)) + { + existingHttpsConnection = new ConcurrentDictionary(); + if (!_existingHttpsConnections.TryAdd(server, existingHttpsConnection)) + throw new InvalidOperationException("Failed to add the canonical HTTPS connection pool."); + } + + if (!existingHttpsConnection.TryGetValue(proxyKey, out HttpsClientConnection connection)) + { + connection = new HttpsClientConnection(server, proxy); + connection.Pooled = true; + + if (!existingHttpsConnection.TryAdd(proxyKey, connection)) + throw new InvalidOperationException("Failed to add the canonical pooled HTTPS connection."); + } + + if (!connection.TryAcquirePooledLease()) + throw new InvalidOperationException("Failed to acquire the canonical pooled HTTPS connection."); + return connection; - }); + } + finally + { + _httpsConnectionsLock.ExitWriteLock(); + } } case DnsTransportProtocol.Quic: { - ConcurrentDictionary existingQuicConnection = _existingQuicConnections.GetOrAdd(server, delegate (NameServerAddress nameServer) - { - return new ConcurrentDictionary(); - }); + NetProxy proxyKey = proxy ?? NetProxy.NONE; - NetProxy proxyKey = proxy; + _quicConnectionsLock.EnterReadLock(); + try + { + if (_existingQuicConnections.TryGetValue(server, out ConcurrentDictionary existingQuicConnection) && existingQuicConnection.TryGetValue(proxyKey, out QuicClientConnection connection)) + { + if (!connection.TryAcquirePooledLease()) + throw new InvalidOperationException("Failed to acquire the canonical pooled QUIC connection."); - if (proxyKey is null) - proxyKey = NetProxy.NONE; + return connection; + } + } + finally + { + _quicConnectionsLock.ExitReadLock(); + } - return existingQuicConnection.GetOrAdd(proxyKey, delegate (NetProxy netProxyKey) + _quicConnectionsLock.EnterWriteLock(); + try { - QuicClientConnection connection = new QuicClientConnection(server, proxy); - connection.Pooled = true; + if (!_existingQuicConnections.TryGetValue(server, out ConcurrentDictionary existingQuicConnection)) + { + existingQuicConnection = new ConcurrentDictionary(); + if (!_existingQuicConnections.TryAdd(server, existingQuicConnection)) + throw new InvalidOperationException("Failed to add the canonical QUIC connection pool."); + } + + if (!existingQuicConnection.TryGetValue(proxyKey, out QuicClientConnection connection)) + { + connection = new QuicClientConnection(server, proxy); + connection.Pooled = true; + + if (!existingQuicConnection.TryAdd(proxyKey, connection)) + throw new InvalidOperationException("Failed to add the canonical pooled QUIC connection."); + } + + if (!connection.TryAcquirePooledLease()) + throw new InvalidOperationException("Failed to acquire the canonical pooled QUIC connection."); + return connection; - }); + } + finally + { + _quicConnectionsLock.ExitWriteLock(); + } } default: diff --git a/TechnitiumLibrary.Net/Dns/ClientConnection/QuicClientConnection.cs b/TechnitiumLibrary.Net/Dns/ClientConnection/QuicClientConnection.cs index c5a6c802..b42e0539 100644 --- a/TechnitiumLibrary.Net/Dns/ClientConnection/QuicClientConnection.cs +++ b/TechnitiumLibrary.Net/Dns/ClientConnection/QuicClientConnection.cs @@ -18,6 +18,7 @@ You should have received a copy of the GNU General Public License */ using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.Diagnostics; using System.IO; @@ -25,6 +26,7 @@ You should have received a copy of the GNU General Public License using System.Net.Quic; using System.Net.Security; using System.Net.Sockets; +using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; using TechnitiumLibrary.Net.Dns.ResourceRecords; @@ -83,6 +85,19 @@ public class QuicClientConnection : DnsClientConnection DateTime _lastQueried; readonly SemaphoreSlim _connectionSemaphore = new SemaphoreSlim(1, 1); + readonly ConditionalWeakTable _streamCapacitySemaphores = new ConditionalWeakTable(); + readonly ConditionalWeakTable _connectionRetirementTasks = new ConditionalWeakTable(); + readonly ConcurrentDictionary _lateTasks = new ConcurrentDictionary(); + readonly ConcurrentDictionary _retirementTasks = new ConcurrentDictionary(); + readonly ConcurrentQueue _retirementFailures = new ConcurrentQueue(); + readonly CancellationTokenSource _stoppingCancellationTokenSource = new CancellationTokenSource(); + readonly TaskCompletionSource _activeQueriesDrained = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + int _stopping; + int _queryAdmissionState; + + const int QUERY_ADMISSION_CLOSED = int.MinValue; + const int ACTIVE_QUERY_COUNT_MASK = int.MaxValue; #endregion @@ -102,51 +117,490 @@ public QuicClientConnection(NameServerAddress server, NetProxy proxy) protected override void Dispose(bool disposing) { if (disposing && !_pooled) + CompleteDisposeAsync(StopAndCloseQueryAdmission()).GetAwaiter().GetResult(); + } + + protected override async ValueTask DisposeAsyncCore() + { + if (!_pooled) + await CompleteDisposeAsync(StopAndCloseQueryAdmission()).ConfigureAwait(false); + } + + private async Task CompleteDisposeAsync(Task activeQueriesDrained) + { + try { - if (_quicConnection is not null) + try + { + await DisposeConnectionAsync(Interlocked.Exchange(ref _quicConnection, null)).ConfigureAwait(false); + } + finally { - _quicConnection.CloseAsync(0).Sync(); - _quicConnection.DisposeAsync().Sync(); + try + { + await activeQueriesDrained.ConfigureAwait(false); + } + finally + { + try + { + await DrainLateTasksAsync().ConfigureAwait(false); + } + finally + { + try + { + await DrainRetirementTasksAsync().ConfigureAwait(false); + } + finally + { + await DisposeConnectionAsync(Interlocked.Exchange(ref _quicConnection, null)).ConfigureAwait(false); + } + } + } } + } + finally + { + Interlocked.Exchange(ref _udpTunnelProxy, null)?.Dispose(); + } + } - _udpTunnelProxy?.Dispose(); + #endregion - _connectionSemaphore?.Dispose(); + #region private + + private static async ValueTask DisposeConnectionAsync(QuicConnection quicConnection) + { + if (quicConnection is null) + return; + + try + { + await quicConnection.CloseAsync(0).ConfigureAwait(false); + } + finally + { + await quicConnection.DisposeAsync().ConfigureAwait(false); } } - protected override async ValueTask DisposeAsyncCore() + private void StartQuery() { - if (!_pooled) + while (true) + { + int admissionState = Volatile.Read(ref _queryAdmissionState); + + if ((admissionState & QUERY_ADMISSION_CLOSED) != 0) + throw new ObjectDisposedException(nameof(QuicClientConnection)); + + if ((admissionState & ACTIVE_QUERY_COUNT_MASK) == ACTIVE_QUERY_COUNT_MASK) + throw new InvalidOperationException("The active QUIC query count overflowed."); + + if (Interlocked.CompareExchange(ref _queryAdmissionState, admissionState + 1, admissionState) == admissionState) + return; + } + } + + private void FinishQuery() + { + while (true) + { + int admissionState = Volatile.Read(ref _queryAdmissionState); + int activeQueryCount = admissionState & ACTIVE_QUERY_COUNT_MASK; + + if (activeQueryCount == 0) + throw new InvalidOperationException("The active QUIC query count is invalid."); + + int newAdmissionState = admissionState - 1; + if (Interlocked.CompareExchange(ref _queryAdmissionState, newAdmissionState, admissionState) != admissionState) + continue; + + if (((newAdmissionState & QUERY_ADMISSION_CLOSED) != 0) && ((newAdmissionState & ACTIVE_QUERY_COUNT_MASK) == 0)) + _activeQueriesDrained.TrySetResult(); + + return; + } + } + + private Task StopAndCloseQueryAdmission() + { + Interlocked.Exchange(ref _stopping, 1); + + try + { + _stoppingCancellationTokenSource.Cancel(); + } + catch (AggregateException ex) + { + Debug.WriteLine(ex); + } + + while (true) + { + int admissionState = Volatile.Read(ref _queryAdmissionState); + + if ((admissionState & QUERY_ADMISSION_CLOSED) != 0) + break; + + int closedAdmissionState = admissionState | QUERY_ADMISSION_CLOSED; + if (Interlocked.CompareExchange(ref _queryAdmissionState, closedAdmissionState, admissionState) == admissionState) + break; + } + + if ((Volatile.Read(ref _queryAdmissionState) & ACTIVE_QUERY_COUNT_MASK) == 0) + _activeQueriesDrained.TrySetResult(); + + return _activeQueriesDrained.Task; + } + + private static async Task DisposeInvalidatedConnectionAsync(QuicConnection quicConnection, UdpTunnelProxy udpTunnelProxy) + { + try + { + await quicConnection.DisposeAsync().ConfigureAwait(false); + } + finally + { + udpTunnelProxy?.Dispose(); + } + } + + private static async Task CompleteInvalidatedConnectionAsync(QuicConnection quicConnection, UdpTunnelProxy udpTunnelProxy, TaskCompletionSource retirementCompleted) + { + try + { + await DisposeInvalidatedConnectionAsync(quicConnection, udpTunnelProxy).ConfigureAwait(false); + retirementCompleted.TrySetResult(); + } + catch (Exception ex) + { + retirementCompleted.TrySetException(ex); + } + } + + private async ValueTask AwaitRetirementAsync(Task retirementTask, long deadline, CancellationToken cancellationToken, CancellationToken waitCancellationToken) + { + bool retirementCompleted = await WaitForTaskAsync(retirementTask, GetRemainingTimeout(deadline), waitCancellationToken).ConfigureAwait(false); + + if (cancellationToken.IsCancellationRequested) + { + TrackRetirementTask(retirementTask); + cancellationToken.ThrowIfCancellationRequested(); + } + + if (!retirementCompleted) + { + TrackRetirementTask(retirementTask); + return false; + } + + await retirementTask.ConfigureAwait(false); + return true; + } + + private async ValueTask InvalidateConnectionAsync(QuicConnection quicConnection, int timeout, CancellationToken cancellationToken) + { + long invalidationDeadline = GetQueryDeadline(timeout, 1); + using CancellationTokenSource invalidationCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _stoppingCancellationTokenSource.Token); + + if (!ReferenceEquals(Volatile.Read(ref _quicConnection), quicConnection)) + { + if (_connectionRetirementTasks.TryGetValue(quicConnection, out Task existingRetirementTask)) + return await AwaitRetirementAsync(existingRetirementTask, invalidationDeadline, cancellationToken, invalidationCancellationTokenSource.Token).ConfigureAwait(false); + + return true; + } + + bool semaphoreAcquired; + + try + { + semaphoreAcquired = await _connectionSemaphore.WaitAsync(GetRemainingTimeout(invalidationDeadline), invalidationCancellationTokenSource.Token).ConfigureAwait(false); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + cancellationToken.ThrowIfCancellationRequested(); + throw; + } + + if (!semaphoreAcquired) + return false; + + Task retirementTask; + + try + { + if (!ReferenceEquals(Volatile.Read(ref _quicConnection), quicConnection)) + { + if (!_connectionRetirementTasks.TryGetValue(quicConnection, out retirementTask)) + return true; + } + else + { + TaskCompletionSource retirementCompleted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + retirementTask = retirementCompleted.Task; + _connectionRetirementTasks.Add(quicConnection, retirementTask); + + Volatile.Write(ref _quicConnection, null); + UdpTunnelProxy udpTunnelProxy = Interlocked.Exchange(ref _udpTunnelProxy, null); + _ = CompleteInvalidatedConnectionAsync(quicConnection, udpTunnelProxy, retirementCompleted); + } + } + finally + { + _connectionSemaphore.Release(); + } + + return await AwaitRetirementAsync(retirementTask, invalidationDeadline, cancellationToken, invalidationCancellationTokenSource.Token).ConfigureAwait(false); + } + + private async Task ObserveRetirementTaskAsync(Task task) + { + try + { + await task.ConfigureAwait(false); + } + catch (Exception ex) + { + _retirementFailures.Enqueue(ex); + ReportRetirementFailure(ex); + } + finally + { + _retirementTasks.TryRemove(task, out _); + } + } + + private void TrackRetirementTask(Task task) + { + if (!_retirementTasks.TryAdd(task, 0)) + return; + + _ = ObserveRetirementTaskAsync(task); + } + + private async Task DrainRetirementTasksAsync() + { + while (!_retirementTasks.IsEmpty) { - if (_quicConnection is not null) + foreach (Task retirementTask in _retirementTasks.Keys) { - await _quicConnection.CloseAsync(0); - await _quicConnection.DisposeAsync(); + try + { + await retirementTask.ConfigureAwait(false); + } + catch + { } } + } + + if (!_retirementFailures.IsEmpty) + { + List retirementFailures = new List(); - _udpTunnelProxy?.Dispose(); + while (_retirementFailures.TryDequeue(out Exception retirementFailure)) + retirementFailures.Add(retirementFailure); - _connectionSemaphore?.Dispose(); + throw new AggregateException("One or more QUIC connection retirements failed.", retirementFailures); } } - #endregion + private static void ReportRetirementFailure(Exception exception) + { + try + { + Console.Error.WriteLine("QUIC connection retirement failed: " + exception); + } + catch + { } + } - #region private + private async Task ObserveLateTaskAsync(Task task) + { + try + { + await task; + } + catch (OperationCanceledException) + { } + catch (ObjectDisposedException) + { } + catch (QuicException) + { } + catch (Exception ex) + { + Debug.WriteLine(ex); + } + finally + { + if (_lateTasks.TryRemove(task, out CancellationTokenSource cancellationTokenSource)) + cancellationTokenSource.Dispose(); + } + } + + private void TrackLateTask(Task task, CancellationTokenSource cancellationTokenSource) + { + if (!_lateTasks.TryAdd(task, cancellationTokenSource)) + throw new InvalidOperationException("The late QUIC task is already tracked."); + + _ = ObserveLateTaskAsync(task); + } + + private async Task DrainLateTasksAsync() + { + while (!_lateTasks.IsEmpty) + { + foreach (KeyValuePair lateTask in _lateTasks) + { + try + { + await lateTask.Key; + } + catch (OperationCanceledException) + { } + catch (ObjectDisposedException) + { } + catch (QuicException) + { } + catch (Exception ex) + { + Debug.WriteLine(ex); + } + + if (_lateTasks.TryRemove(lateTask.Key, out CancellationTokenSource cancellationTokenSource)) + cancellationTokenSource.Dispose(); + } + } + } + + private static async Task ObserveTaskAsync(Task task) + { + try + { + await task; + } + catch (OperationCanceledException) + { } + catch (ObjectDisposedException) + { } + catch (QuicException) + { } + catch (Exception ex) + { + Debug.WriteLine(ex); + } + } + + private static async Task WaitForTaskAsync(Task task, int timeout, CancellationToken cancellationToken) + { + using (CancellationTokenSource delayCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken)) + { + try + { + Task completedTask = await Task.WhenAny(task, Task.Delay(timeout, delayCancellationTokenSource.Token)); + return completedTask == task; + } + finally + { + delayCancellationTokenSource.Cancel(); + } + } + } + + private static long GetQueryDeadline(int timeout, long attempts) + { + if (timeout == Timeout.Infinite) + return long.MaxValue; + + long now = Environment.TickCount64; + attempts = Math.Max(attempts, 1L); + + if (timeout <= 0) + return now; + + if (attempts > (long.MaxValue - now) / timeout) + return long.MaxValue; + + return now + (timeout * attempts); + } + + private static int GetRemainingTimeout(long deadline) + { + if (deadline == long.MaxValue) + return Timeout.Infinite; + + long remaining = deadline - Environment.TickCount64; + + if (remaining <= 0) + return 0; + + return remaining > int.MaxValue ? int.MaxValue : Convert.ToInt32(remaining); + } + + private static int GetShorterTimeout(int firstTimeout, int secondTimeout) + { + if (firstTimeout == Timeout.Infinite) + return secondTimeout; + + if (secondTimeout == Timeout.Infinite) + return firstTimeout; + + return Math.Min(firstTimeout, secondTimeout); + } + + private static void AbortCancelledStream(object state) + { + try + { + ((QuicStream)state).Abort(QuicAbortDirection.Both, (long)DnsOverQuicErrorCodes.DOQ_REQUEST_CANCELLED); + } + catch (ObjectDisposedException) + { } + catch (QuicException) + { } + } + + private SemaphoreSlim GetStreamCapacitySemaphore(QuicConnection quicConnection) + { + return _streamCapacitySemaphores.GetValue(quicConnection, static delegate { return new SemaphoreSlim(0); }); + } + + private void OnStreamCapacityChanged(QuicConnection quicConnection, QuicStreamCapacityChangedArgs args) + { + if (args.BidirectionalIncrement <= 0) + return; + + try + { + GetStreamCapacitySemaphore(quicConnection).Release(args.BidirectionalIncrement); + } + catch (SemaphoreFullException ex) + { + Debug.WriteLine(ex); + } + } private async Task GetConnectionAsync(int timeout, CancellationToken cancellationToken) { - if (_quicConnection is not null) - return _quicConnection; + if (Volatile.Read(ref _stopping) != 0) + throw new ObjectDisposedException(nameof(QuicClientConnection)); + + QuicConnection existingConnection = Volatile.Read(ref _quicConnection); + if (existingConnection is not null) + return existingConnection; if (!await _connectionSemaphore.WaitAsync(timeout, cancellationToken)) return null; //timed out try { - if (_quicConnection is not null) - return _quicConnection; + existingConnection = Volatile.Read(ref _quicConnection); + if (existingConnection is not null) + return existingConnection; + + if (Volatile.Read(ref _stopping) != 0) + throw new ObjectDisposedException(nameof(QuicClientConnection)); IPEndPoint remoteEP; @@ -175,6 +629,7 @@ private async Task GetConnectionAsync(int timeout, CancellationT DefaultStreamErrorCode = (long)DnsOverQuicErrorCodes.DOQ_REQUEST_CANCELLED, MaxInboundUnidirectionalStreams = 0, MaxInboundBidirectionalStreams = 0, + StreamCapacityCallback = OnStreamCapacityChanged, ClientAuthenticationOptions = new SslClientAuthenticationOptions { ApplicationProtocols = new List() { new SslApplicationProtocol("doq") }, @@ -202,23 +657,68 @@ private async Task GetConnectionAsync(int timeout, CancellationT } } - _quicConnection = await TaskExtensions.TimeoutAsync(async delegate (CancellationToken cancellationToken1) + using CancellationTokenSource connectCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); + connectCancellationTokenSource.CancelAfter(30000); + + QuicConnection newConnection; + + try + { + newConnection = await QuicConnection.ConnectAsync(connectionOptions, connectCancellationTokenSource.Token).ConfigureAwait(false); + } + catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested && connectCancellationTokenSource.IsCancellationRequested) + { + throw new TimeoutException(); + } + + if (Volatile.Read(ref _stopping) != 0) + { + await DisposeConnectionAsync(newConnection).ConfigureAwait(false); + throw new ObjectDisposedException(nameof(QuicClientConnection)); + } + + Volatile.Write(ref _quicConnection, newConnection); + + if (Volatile.Read(ref _stopping) != 0) { - return await QuicConnection.ConnectAsync(connectionOptions, cancellationToken1); - }, 30000, cancellationToken); + if (ReferenceEquals(Interlocked.CompareExchange(ref _quicConnection, null, newConnection), newConnection)) + await DisposeConnectionAsync(newConnection).ConfigureAwait(false); - return _quicConnection; + throw new ObjectDisposedException(nameof(QuicClientConnection)); + } + + return newConnection; } finally { _connectionSemaphore.Release(); + + if (Volatile.Read(ref _stopping) != 0) + Interlocked.Exchange(ref _udpTunnelProxy, null)?.Dispose(); } } - private static async Task QuicQueryAsync(DnsDatagram request, QuicConnection quicConnection, CancellationToken cancellationToken) + private async Task QuicQueryAsync(DnsDatagram request, QuicConnection quicConnection, CancellationToken cancellationToken) { - await using (QuicStream quicStream = await quicConnection.OpenOutboundStreamAsync(QuicStreamType.Bidirectional, cancellationToken)) + SemaphoreSlim streamCapacitySemaphore = GetStreamCapacitySemaphore(quicConnection); + await streamCapacitySemaphore.WaitAsync(cancellationToken); + + try + { + cancellationToken.ThrowIfCancellationRequested(); + } + catch + { + streamCapacitySemaphore.Release(); + throw; + } + + await using (QuicStream quicStream = await quicConnection.OpenOutboundStreamAsync(QuicStreamType.Bidirectional, CancellationToken.None)) { + await using CancellationTokenRegistration cancellationTokenRegistration = cancellationToken.Register(AbortCancelledStream, quicStream); + + cancellationToken.ThrowIfCancellationRequested(); + //serialize and send request with FIN flag using (MemoryStream mS = new MemoryStream(64)) { @@ -280,6 +780,31 @@ private static async Task QuicQueryAsync(DnsDatagram request, QuicC #region public public override async Task QueryAsync(DnsDatagram request, int timeout, int retries, CancellationToken cancellationToken) + { + StartQuery(); + + try + { + Task queryTask; + + if ((SynchronizationContext.Current is null) && (TaskScheduler.Current == TaskScheduler.Default)) + { + queryTask = QueryCoreAsync(request, timeout, retries, cancellationToken); + } + else + { + queryTask = Task.Run(delegate { return QueryCoreAsync(request, timeout, retries, cancellationToken); }); + } + + return await queryTask.ConfigureAwait(false); + } + finally + { + FinishQuery(); + } + } + + private async Task QueryCoreAsync(DnsDatagram request, int timeout, int retries, CancellationToken cancellationToken) { _lastQueried = DateTime.UtcNow; @@ -287,6 +812,9 @@ public override async Task QueryAsync(DnsDatagram request, int time stopwatch.Start(); + long requestDeadline = GetQueryDeadline(timeout, Math.Max((long)retries, 1L) * 2L); + bool queryDeadlineSet = false; + long queryDeadline = 0; int retry = 0; while (retry < retries) //retry loop { @@ -294,91 +822,222 @@ public override async Task QueryAsync(DnsDatagram request, int time retry++; - Task quicConnectionTask = GetConnectionAsync(timeout, cancellationToken); + int remainingRequestTimeout = GetRemainingTimeout(requestDeadline); + if (remainingRequestTimeout == 0) + break; + + if (queryDeadlineSet) + { + int remainingQueryTimeout = GetRemainingTimeout(queryDeadline); + if (remainingQueryTimeout == 0) + break; + + remainingRequestTimeout = GetShorterTimeout(remainingRequestTimeout, remainingQueryTimeout); + } + + int connectionTimeout = GetShorterTimeout(timeout, remainingRequestTimeout); + CancellationTokenSource connectionCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _stoppingCancellationTokenSource.Token); + bool connectionCancellationOwnershipTransferred = false; + Task quicConnectionTask = null; + QuicConnection quicConnection; //wait for connection with timeout - using (CancellationTokenSource timeoutCancellationTokenSource = new CancellationTokenSource()) + try { - await using (CancellationTokenRegistration ctr = cancellationToken.Register(timeoutCancellationTokenSource.Cancel)) + quicConnectionTask = GetConnectionAsync(connectionTimeout, connectionCancellationTokenSource.Token); + + if (!await WaitForTaskAsync(quicConnectionTask, connectionTimeout, cancellationToken)) { - if ((await Task.WhenAny(quicConnectionTask, Task.Delay(timeout, timeoutCancellationTokenSource.Token)) != quicConnectionTask) && (quicConnectionTask.Status != TaskStatus.RanToCompletion)) - continue; //request timed out; retry + connectionCancellationTokenSource.Cancel(); + bool connectionCleanupCompleted = quicConnectionTask.IsCompleted; + + remainingRequestTimeout = GetRemainingTimeout(requestDeadline); + if (queryDeadlineSet) + remainingRequestTimeout = GetShorterTimeout(remainingRequestTimeout, GetRemainingTimeout(queryDeadline)); + + if (!connectionCleanupCompleted && (remainingRequestTimeout != 0)) + connectionCleanupCompleted = await WaitForTaskAsync(quicConnectionTask, remainingRequestTimeout, cancellationToken); + + if (cancellationToken.IsCancellationRequested) + { + if (connectionCleanupCompleted) + { + await ObserveTaskAsync(quicConnectionTask); + } + else + { + TrackLateTask(quicConnectionTask, connectionCancellationTokenSource); + connectionCancellationOwnershipTransferred = true; + } + + cancellationToken.ThrowIfCancellationRequested(); + } + + if (!connectionCleanupCompleted) + { + TrackLateTask(quicConnectionTask, connectionCancellationTokenSource); + connectionCancellationOwnershipTransferred = true; + break; + } + + await ObserveTaskAsync(quicConnectionTask); + continue; } - timeoutCancellationTokenSource.Cancel(); //to stop delay task + try + { + quicConnection = await quicConnectionTask; + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + cancellationToken.ThrowIfCancellationRequested(); + throw; + } + } + finally + { + if (!connectionCancellationOwnershipTransferred) + connectionCancellationTokenSource.Dispose(); } - QuicConnection quicConnection = await quicConnectionTask; + cancellationToken.ThrowIfCancellationRequested(); + if (quicConnection is null) continue; //semaphone wait timed out; retry - Task task; + if (!queryDeadlineSet) + { + queryDeadline = GetQueryDeadline(timeout, retries - retry + 1); + queryDeadlineSet = true; + } + + int remainingTimeout = GetRemainingTimeout(queryDeadline); + if (remainingTimeout == 0) + break; + + int attemptTimeout = timeout == Timeout.Infinite ? Timeout.Infinite : Math.Min(timeout, remainingTimeout); + CancellationTokenSource queryCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _stoppingCancellationTokenSource.Token); + bool queryCancellationOwnershipTransferred = false; + Task task = null; //query and wait for response with timeout - using (CancellationTokenSource timeoutCancellationTokenSource = new CancellationTokenSource()) + try { - await using (CancellationTokenRegistration ctr = cancellationToken.Register(timeoutCancellationTokenSource.Cancel)) + task = QuicQueryAsync(request, quicConnection, queryCancellationTokenSource.Token); + + if (!await WaitForTaskAsync(task, attemptTimeout, cancellationToken)) { - task = QuicQueryAsync(request, quicConnection, timeoutCancellationTokenSource.Token); + queryCancellationTokenSource.Cancel(); - if ((await Task.WhenAny(task, Task.Delay(timeout, timeoutCancellationTokenSource.Token)) != task) && (task.Status != TaskStatus.RanToCompletion)) + if (cancellationToken.IsCancellationRequested) { - timeoutCancellationTokenSource.Cancel(); //to stop running task - continue; //request timed out; retry + if (task.IsCompleted) + { + await ObserveTaskAsync(task); + } + else + { + TrackLateTask(task, queryCancellationTokenSource); + queryCancellationOwnershipTransferred = true; + } + + cancellationToken.ThrowIfCancellationRequested(); } - } - timeoutCancellationTokenSource.Cancel(); //to stop delay task - } + remainingTimeout = GetRemainingTimeout(queryDeadline); + bool queryCleanupCompleted = task.IsCompleted; - DnsDatagram response; + if (!queryCleanupCompleted && (remainingTimeout != 0)) + queryCleanupCompleted = await WaitForTaskAsync(task, remainingTimeout, cancellationToken); - try - { - response = await task; - } - catch (ObjectDisposedException) - { - //ensure existing connection is disposed to allow reconnection later - await quicConnection.DisposeAsync(); - _quicConnection = null; - _udpTunnelProxy?.Dispose(); + if (cancellationToken.IsCancellationRequested) + { + if (queryCleanupCompleted) + { + await ObserveTaskAsync(task); + } + else + { + TrackLateTask(task, queryCancellationTokenSource); + queryCancellationOwnershipTransferred = true; + } + + cancellationToken.ThrowIfCancellationRequested(); + } - if (retry == 1) - { - //quic connection was disposed on first attempt; retry to reconnect - retry = 0; + if (!queryCleanupCompleted) + { + TrackLateTask(task, queryCancellationTokenSource); + queryCancellationOwnershipTransferred = true; + break; + } + + await ObserveTaskAsync(task); continue; } - throw; - } - catch (QuicException ex) - { - //close existing connection to allow reconnection later - await quicConnection.DisposeAsync(); - _quicConnection = null; - _udpTunnelProxy?.Dispose(); + DnsDatagram response; - if (((ex.QuicError == QuicError.ConnectionIdle) || (ex.QuicError == QuicError.ConnectionAborted)) && (retry == 1)) + try { - //connection idle/aborted on first attempt; retry to reconnect - retry = 0; - continue; + response = await task; + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + cancellationToken.ThrowIfCancellationRequested(); + throw; } + catch (ObjectDisposedException) + { + //ensure existing connection is disposed to allow reconnection later + int invalidationTimeout = GetShorterTimeout(GetRemainingTimeout(requestDeadline), GetRemainingTimeout(queryDeadline)); + if (!await InvalidateConnectionAsync(quicConnection, invalidationTimeout, cancellationToken).ConfigureAwait(false)) + break; - throw; - } + if (retry == 1) + { + //quic connection was disposed on first attempt; retry to reconnect + retry = 0; + continue; + } + + throw; + } + catch (QuicException ex) + { + //close existing connection to allow reconnection later + int invalidationTimeout = GetShorterTimeout(GetRemainingTimeout(requestDeadline), GetRemainingTimeout(queryDeadline)); + if (!await InvalidateConnectionAsync(quicConnection, invalidationTimeout, cancellationToken).ConfigureAwait(false)) + break; + + if (((ex.QuicError == QuicError.ConnectionIdle) || (ex.QuicError == QuicError.ConnectionAborted)) && (retry == 1)) + { + //connection idle/aborted on first attempt; retry to reconnect + retry = 0; + continue; + } - stopwatch.Stop(); + throw; + } - response.SetMetadata(_server, stopwatch.Elapsed.TotalMilliseconds); + cancellationToken.ThrowIfCancellationRequested(); - ValidateResponse(request, response); + stopwatch.Stop(); - return response; + response.SetMetadata(_server, stopwatch.Elapsed.TotalMilliseconds); + + ValidateResponse(request, response); + + return response; + } + finally + { + if (!queryCancellationOwnershipTransferred) + queryCancellationTokenSource.Dispose(); + } } + cancellationToken.ThrowIfCancellationRequested(); throw new DnsClientNoResponseException("DnsClient failed to resolve the request" + (request.Question.Count > 0 ? " '" + request.Question[0].ToString() + "'" : "") + ": request timed out for name server [" + _server.ToString() + "]."); }