Skip to content

Commit 2f96e88

Browse files
ManickaPManickaP
authored
[QUIC] Call silent shutdown in case CloseAsync failed. (#96807)
* Cancelletion is temporary * Use ordinary TCS and revert changes in ValueTaskSource * Lower probability of creating extra TCS instance. * Revert "Lower probability of creating extra TCS instance." This reverts commit 636037cf9e85bc1f275d6cb7346866d4b86c5938. * Revert "Use ordinary TCS and revert changes in ValueTaskSource" This reverts commit f393313f886bf2073ddaea9b02f9bb30b80f047b. * Added Radek's test * Use RVTS * Fix the fix :D
1 parent 328ef20 commit 2f96e88

File tree

5 files changed

+85
-19
lines changed

5 files changed

+85
-19
lines changed

‎src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/ResettableValueTaskSource.cs

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -105,9 +105,7 @@ public bool TryGetValueTask(out ValueTask valueTask, object? keepAlive = null, C
105105
_state = State.Awaiting;
106106
}
107107
// None, Ready, Completed: return the current task.
108-
if (state == State.None ||
109-
state == State.Ready ||
110-
state == State.Completed)
108+
if (state is State.None or State.Ready or State.Completed)
111109
{
112110
// Remember that the value task with the current version is being given out.
113111
_hasWaiter = true;
@@ -167,8 +165,7 @@ private bool TryComplete(Exception? exception, bool final)
167165

168166
// If the _valueTaskSource has already been set, we don't want to lose the result by overwriting it.
169167
// So keep it as is and store the result in _finalTaskSource.
170-
if (state == State.None ||
171-
state == State.Awaiting)
168+
if (state is State.None or State.Awaiting)
172169
{
173170
_state = final ? State.Completed : State.Ready;
174171
}
@@ -178,16 +175,14 @@ private bool TryComplete(Exception? exception, bool final)
178175
{
179176
// Set up the exception stack trace for the caller.
180177
exception = exception.StackTrace is null ? ExceptionDispatchInfo.SetCurrentStackTrace(exception) : exception;
181-
if (state == State.None ||
182-
state == State.Awaiting)
178+
if (state is State.None or State.Awaiting)
183179
{
184180
_valueTaskSource.SetException(exception);
185181
}
186182
}
187183
else
188184
{
189-
if (state == State.None ||
190-
state == State.Awaiting)
185+
if (state is State.None or State.Awaiting)
191186
{
192187
_valueTaskSource.SetResult(final);
193188
}

‎src/libraries/System.Net.Quic/src/System/Net/Quic/QuicConnection.cs

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,26 @@ static async ValueTask<QuicConnection> StartConnectAsync(QuicClientConnectionOpt
109109
private int _disposed;
110110

111111
private readonly ValueTaskSource _connectedTcs = new ValueTaskSource();
112-
private readonly ValueTaskSource _shutdownTcs = new ValueTaskSource();
112+
private readonly ResettableValueTaskSource _shutdownTcs = new ResettableValueTaskSource()
113+
{
114+
CancellationAction = target =>
115+
{
116+
try
117+
{
118+
if (target is QuicConnection connection)
119+
{
120+
// The OCE will be propagated through stored CancellationToken in ResettableValueTaskSource.
121+
connection._shutdownTcs.TrySetResult();
122+
}
123+
}
124+
catch (ObjectDisposedException)
125+
{
126+
// We collided with a Dispose in another thread. This can happen
127+
// when using CancellationTokenSource.CancelAfter.
128+
// Ignore the exception
129+
}
130+
}
131+
};
113132

114133
private readonly CancellationTokenSource _shutdownTokenSource = new CancellationTokenSource();
115134

@@ -467,7 +486,7 @@ public ValueTask CloseAsync(long errorCode, CancellationToken cancellationToken
467486
{
468487
ObjectDisposedException.ThrowIf(_disposed == 1, this);
469488

470-
if (_shutdownTcs.TryInitialize(out ValueTask valueTask, this, cancellationToken))
489+
if (_shutdownTcs.TryGetValueTask(out ValueTask valueTask, this, cancellationToken))
471490
{
472491
unsafe
473492
{
@@ -520,7 +539,7 @@ private unsafe int HandleEventShutdownComplete()
520539
_acceptQueue.Writer.TryComplete(exception);
521540
_connectedTcs.TrySetException(exception);
522541
_shutdownTokenSource.Cancel();
523-
_shutdownTcs.TrySetResult();
542+
_shutdownTcs.TrySetResult(final: true);
524543
return QUIC_STATUS_SUCCESS;
525544
}
526545
private unsafe int HandleEventLocalAddressChanged(ref LOCAL_ADDRESS_CHANGED_DATA data)
@@ -626,7 +645,7 @@ public async ValueTask DisposeAsync()
626645
}
627646

628647
// Check if the connection has been shut down and if not, shut it down.
629-
if (_shutdownTcs.TryInitialize(out ValueTask valueTask, this))
648+
if (_shutdownTcs.TryGetValueTask(out ValueTask valueTask, this))
630649
{
631650
unsafe
632651
{
@@ -636,9 +655,19 @@ public async ValueTask DisposeAsync()
636655
(ulong)_defaultCloseErrorCode);
637656
}
638657
}
658+
else if (!valueTask.IsCompletedSuccessfully)
659+
{
660+
unsafe
661+
{
662+
MsQuicApi.Api.ConnectionShutdown(
663+
_handle,
664+
QUIC_CONNECTION_SHUTDOWN_FLAGS.SILENT,
665+
(ulong)_defaultCloseErrorCode);
666+
}
667+
}
639668

640669
// Wait for SHUTDOWN_COMPLETE, the last event, so that all resources can be safely released.
641-
await valueTask.ConfigureAwait(false);
670+
await _shutdownTcs.GetFinalTask(this).ConfigureAwait(false);
642671
Debug.Assert(_connectedTcs.IsCompleted);
643672
_handle.Dispose();
644673
_shutdownTokenSource.Dispose();

‎src/libraries/System.Net.Quic/src/System/Net/Quic/QuicListener.cs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -220,7 +220,7 @@ private async void StartConnectionHandshake(QuicConnection connection, SslClient
220220
{
221221
using CancellationTokenSource linkedCts = CancellationTokenSource.CreateLinkedTokenSource(_disposeCts.Token, connection.ConnectionShutdownToken);
222222
cancellationToken = linkedCts.Token;
223-
// initial timeout for retrieving connection options
223+
// Initial timeout for retrieving connection options.
224224
linkedCts.CancelAfter(handshakeTimeout);
225225

226226
wrapException = true;
@@ -229,7 +229,7 @@ private async void StartConnectionHandshake(QuicConnection connection, SslClient
229229

230230
options.Validate(nameof(options));
231231

232-
// update handshake timetout based on the returned value
232+
// Update handshake timeout based on the returned value.
233233
handshakeTimeout = options.HandshakeTimeout;
234234
linkedCts.CancelAfter(handshakeTimeout);
235235

@@ -248,12 +248,12 @@ private async void StartConnectionHandshake(QuicConnection connection, SslClient
248248
NetEventSource.Info(connection, $"{connection} Connection closed by remote peer");
249249
}
250250

251-
// retrieve the exception which failed the handshake, the parameters are not going to be
252-
// validated because the inner _connectedTcs is already transitioned to faulted state
251+
// Retrieve the exception which failed the handshake, the parameters are not going to be
252+
// validated because the inner _connectedTcs is already transitioned to faulted state.
253253
ValueTask task = connection.FinishHandshakeAsync(null!, null!, default);
254254
Debug.Assert(task.IsFaulted);
255255

256-
// unwrap AggregateException and propagate it to the accept queue
256+
// Unwrap AggregateException and propagate it to the accept queue.
257257
Exception ex = task.AsTask().Exception!.InnerException!;
258258

259259
await connection.DisposeAsync().ConfigureAwait(false);

‎src/libraries/System.Net.Quic/src/System/Net/Quic/QuicStream.cs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -613,6 +613,7 @@ private unsafe int HandleEventShutdownComplete(ref SHUTDOWN_COMPLETE_DATA data)
613613
_receiveTcs.TrySetException(exception, final: true);
614614
_sendTcs.TrySetException(exception, final: true);
615615
}
616+
_startedTcs.TrySetResult();
616617
_shutdownTcs.TrySetResult();
617618
return QUIC_STATUS_SUCCESS;
618619
}

‎src/libraries/System.Net.Quic/tests/FunctionalTests/QuicConnectionTests.cs

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,47 @@ await RunClientServer(
124124
});
125125
}
126126

127+
[Fact]
128+
public async Task DisposeAfterCloseCanceled()
129+
{
130+
using var sync = new SemaphoreSlim(0);
131+
132+
await RunClientServer(
133+
async clientConnection =>
134+
{
135+
var cts = new CancellationTokenSource();
136+
cts.Cancel();
137+
await Assert.ThrowsAsync<OperationCanceledException>(async () => await clientConnection.CloseAsync(ExpectedErrorCode, cts.Token));
138+
await clientConnection.DisposeAsync();
139+
sync.Release();
140+
},
141+
async serverConnection =>
142+
{
143+
await sync.WaitAsync();
144+
await serverConnection.DisposeAsync();
145+
});
146+
}
147+
148+
[Fact]
149+
public async Task DisposeAfterCloseTaskStored()
150+
{
151+
using var sync = new SemaphoreSlim(0);
152+
153+
await RunClientServer(
154+
async clientConnection =>
155+
{
156+
var cts = new CancellationTokenSource();
157+
var task = clientConnection.CloseAsync(0).AsTask();
158+
await clientConnection.DisposeAsync();
159+
sync.Release();
160+
},
161+
async serverConnection =>
162+
{
163+
await sync.WaitAsync();
164+
await serverConnection.DisposeAsync();
165+
});
166+
}
167+
127168
[Fact]
128169
public async Task ConnectionClosedByPeer_WithPendingAcceptAndConnect_PendingAndSubsequentThrowConnectionAbortedException()
129170
{

0 commit comments

Comments
 (0)