Skip to content

Commit 7d92975

Browse files
JamesNKJamesNK
committed
PR feedback
1 parent ec05ce5 commit 7d92975

File tree

9 files changed

+44
-75
lines changed

9 files changed

+44
-75
lines changed

‎src/Servers/Kestrel/Core/src/Internal/Http3/Http3Stream.cs‎

Lines changed: 18 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -496,25 +496,26 @@ public async Task ProcessRequestAsync<TContext>(IHttpApplication<TContext> appli
496496
}
497497
finally
498498
{
499-
// Completing the stream needs to happen in the middle of DisposeAsync.
500-
// After the stream has drained but before the stream is returned to the pool.
501-
// OnCompleted callback executes in the right place.
502-
//
503-
// TODO consider adding new feature to remove need for this:
504-
// https://github.com/dotnet/aspnetcore/issues/35400
505-
_context.StreamContext.Features.Get<IConnectionCompleteFeature>()!.OnCompleted(static (state) =>
506-
{
507-
var s = (Http3Stream)state;
508-
509-
s.ApplyCompletionFlag(StreamCompletionFlags.Completed);
510-
511-
// Tells the connection to remove the stream from its active collection.
512-
s._context.StreamLifetimeHandler.OnStreamCompleted(s);
499+
// Drain transports and dispose.
500+
await _context.StreamContext.DisposeAsync();
513501

514-
return Task.CompletedTask;
515-
}, this);
502+
// Tells the connection to remove the stream from its active collection.
503+
ApplyCompletionFlag(StreamCompletionFlags.Completed);
504+
_context.StreamLifetimeHandler.OnStreamCompleted(this);
516505

517-
await _context.StreamContext.DisposeAsync();
506+
// TODO this is a hack for .NET 6 pooling.
507+
//
508+
// Pooling needs to happen after transports have been drained and stream
509+
// has been completed and is no longer active. All of this logic can't
510+
// be placed in ConnectionContext.DisposeAsync. Instead, QuicStreamContext
511+
// has pooling happen in QuicStreamContext.Dispose.
512+
//
513+
// ConnectionContext only implements IDisposableAsync by default. Only
514+
// QuicStreamContext should pass this check.
515+
if (_context.StreamContext is IDisposable disposableStream)
516+
{
517+
disposableStream.Dispose();
518+
}
518519
}
519520
}
520521
}

‎src/Servers/Kestrel/Transport.Quic/src/Internal/QuicStreamContext.FeatureCollection.cs‎

Lines changed: 1 addition & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,14 @@
11
// Licensed to the .NET Foundation under one or more agreements.
22
// The .NET Foundation licenses this file to you under the MIT license.
33

4-
using System.Net.Sockets;
54
using Microsoft.AspNetCore.Connections;
65
using Microsoft.AspNetCore.Connections.Features;
7-
using Microsoft.Extensions.Logging;
86

97
namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Quic.Internal
108
{
11-
internal sealed partial class QuicStreamContext : IPersistentStateFeature, IStreamDirectionFeature, IProtocolErrorCodeFeature, IStreamIdFeature, IStreamAbortFeature, IConnectionCompleteFeature
9+
internal sealed partial class QuicStreamContext : IPersistentStateFeature, IStreamDirectionFeature, IProtocolErrorCodeFeature, IStreamIdFeature, IStreamAbortFeature
1210
{
1311
private IDictionary<object, object?>? _persistentState;
14-
private Stack<KeyValuePair<Func<object, Task>, object>>? _onCompleted;
1512

1613
public bool CanRead { get; private set; }
1714
public bool CanWrite { get; private set; }
@@ -69,23 +66,13 @@ public void AbortWrite(long errorCode, ConnectionAbortedException abortReason)
6966
}
7067
}
7168

72-
public void OnCompleted(Func<object, Task> callback, object state)
73-
{
74-
if (_onCompleted == null)
75-
{
76-
_onCompleted = new Stack<KeyValuePair<Func<object, Task>, object>>();
77-
}
78-
_onCompleted.Push(new KeyValuePair<Func<object, Task>, object>(callback, state));
79-
}
80-
8169
private void InitializeFeatures()
8270
{
8371
_currentIPersistentStateFeature = this;
8472
_currentIStreamDirectionFeature = this;
8573
_currentIProtocolErrorCodeFeature = this;
8674
_currentIStreamIdFeature = this;
8775
_currentIStreamAbortFeature = this;
88-
_currentIConnectionCompleteFeature = this;
8976
}
9077
}
9178
}

‎src/Servers/Kestrel/Transport.Quic/src/Internal/QuicStreamContext.cs‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Quic.Internal
1919
{
20-
internal partial class QuicStreamContext : TransportConnection, IPooledStream
20+
internal partial class QuicStreamContext : TransportConnection, IPooledStream, IDisposable
2121
{
2222
// Internal for testing.
2323
internal Task _processingTask = Task.CompletedTask;
@@ -444,12 +444,18 @@ public override async ValueTask DisposeAsync()
444444

445445
lock (_shutdownLock)
446446
{
447+
if (!CanReuse)
448+
{
449+
DisposeCore();
450+
}
451+
447452
_stream.Dispose();
448453
_stream = null!;
449454
}
455+
}
450456

451-
await ConnectionCompletion.FireOnCompletedAsync(_log, _onCompleted);
452-
457+
public void Dispose()
458+
{
453459
if (!_connection.TryReturnStream(this))
454460
{
455461
// Dispose when one of:

‎src/Servers/Kestrel/Transport.Quic/test/QuicConnectionContextTests.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -578,6 +578,7 @@ static async Task SendStream(RequestState requestState)
578578
// Both send and receive loops have exited.
579579
await quicStreamContext._processingTask.DefaultTimeout();
580580
await quicStreamContext.DisposeAsync();
581+
quicStreamContext.Dispose();
581582
}
582583
}
583584

@@ -614,6 +615,7 @@ public async Task PersistentState_StreamsReused_StatePersisted()
614615
var quicStreamContext1 = Assert.IsType<QuicStreamContext>(serverStream1);
615616
await quicStreamContext1._processingTask.DefaultTimeout();
616617
await quicStreamContext1.DisposeAsync();
618+
quicStreamContext1.Dispose();
617619

618620
var clientStream2 = clientConnection.OpenBidirectionalStream();
619621
await clientStream2.WriteAsync(TestData, endStream: true).DefaultTimeout();
@@ -634,6 +636,7 @@ public async Task PersistentState_StreamsReused_StatePersisted()
634636
var quicStreamContext2 = Assert.IsType<QuicStreamContext>(serverStream2);
635637
await quicStreamContext2._processingTask.DefaultTimeout();
636638
await quicStreamContext2.DisposeAsync();
639+
quicStreamContext2.Dispose();
637640

638641
Assert.Same(quicStreamContext1, quicStreamContext2);
639642

‎src/Servers/Kestrel/Transport.Quic/test/QuicStreamContextTests.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ public async Task BidirectionalStream_ServerWritesDataAndDisposes_ClientReadsDat
101101

102102
Logger.LogInformation("Server disposing stream.");
103103
await quicStreamContext.DisposeAsync().DefaultTimeout();
104+
quicStreamContext.Dispose();
104105

105106
Logger.LogInformation("Client reading until end of stream.");
106107
var data = await clientStream.ReadUntilEndAsync().DefaultTimeout();

‎src/Servers/Kestrel/Transport.Quic/test/QuicTestHelpers.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@ public static async Task<QuicStreamContext> CreateAndCompleteBidirectionalStream
130130
Assert.True(quicStreamContext.CanRead);
131131

132132
await quicStreamContext.DisposeAsync();
133+
quicStreamContext.Dispose();
133134

134135
return quicStreamContext;
135136
}

‎src/Servers/Kestrel/shared/TransportConnection.Generated.cs‎

Lines changed: 0 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@ internal partial class TransportConnection : IFeatureCollection,
3434
internal protected IStreamDirectionFeature? _currentIStreamDirectionFeature;
3535
internal protected IStreamIdFeature? _currentIStreamIdFeature;
3636
internal protected IStreamAbortFeature? _currentIStreamAbortFeature;
37-
internal protected IConnectionCompleteFeature? _currentIConnectionCompleteFeature;
3837

3938
private int _featureRevision;
4039

@@ -54,7 +53,6 @@ private void FastReset()
5453
_currentIStreamDirectionFeature = null;
5554
_currentIStreamIdFeature = null;
5655
_currentIStreamAbortFeature = null;
57-
_currentIConnectionCompleteFeature = null;
5856
}
5957

6058
// Internal for testing
@@ -170,10 +168,6 @@ private void ExtraFeatureSet(Type key, object? value)
170168
{
171169
feature = _currentIStreamAbortFeature;
172170
}
173-
else if (key == typeof(IConnectionCompleteFeature))
174-
{
175-
feature = _currentIConnectionCompleteFeature;
176-
}
177171
else if (MaybeExtra != null)
178172
{
179173
feature = ExtraFeatureGet(key);
@@ -230,10 +224,6 @@ private void ExtraFeatureSet(Type key, object? value)
230224
{
231225
_currentIStreamAbortFeature = (IStreamAbortFeature?)value;
232226
}
233-
else if (key == typeof(IConnectionCompleteFeature))
234-
{
235-
_currentIConnectionCompleteFeature = (IConnectionCompleteFeature?)value;
236-
}
237227
else
238228
{
239229
ExtraFeatureSet(key, value);
@@ -292,10 +282,6 @@ private void ExtraFeatureSet(Type key, object? value)
292282
{
293283
feature = Unsafe.As<IStreamAbortFeature?, TFeature?>(ref _currentIStreamAbortFeature);
294284
}
295-
else if (typeof(TFeature) == typeof(IConnectionCompleteFeature))
296-
{
297-
feature = Unsafe.As<IConnectionCompleteFeature?, TFeature?>(ref _currentIConnectionCompleteFeature);
298-
}
299285
else if (MaybeExtra != null)
300286
{
301287
feature = (TFeature?)(ExtraFeatureGet(typeof(TFeature)));
@@ -360,10 +346,6 @@ private void ExtraFeatureSet(Type key, object? value)
360346
{
361347
_currentIStreamAbortFeature = Unsafe.As<TFeature?, IStreamAbortFeature?>(ref feature);
362348
}
363-
else if (typeof(TFeature) == typeof(IConnectionCompleteFeature))
364-
{
365-
_currentIConnectionCompleteFeature = Unsafe.As<TFeature?, IConnectionCompleteFeature?>(ref feature);
366-
}
367349
else
368350
{
369351
ExtraFeatureSet(typeof(TFeature), feature);
@@ -416,10 +398,6 @@ private IEnumerable<KeyValuePair<Type, object>> FastEnumerable()
416398
{
417399
yield return new KeyValuePair<Type, object>(typeof(IStreamAbortFeature), _currentIStreamAbortFeature);
418400
}
419-
if (_currentIConnectionCompleteFeature != null)
420-
{
421-
yield return new KeyValuePair<Type, object>(typeof(IConnectionCompleteFeature), _currentIConnectionCompleteFeature);
422-
}
423401

424402
if (MaybeExtra != null)
425403
{

‎src/Servers/Kestrel/shared/test/Http3/Http3InMemory.cs‎

Lines changed: 10 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1001,7 +1001,7 @@ public void RequestClose()
10011001
}
10021002
}
10031003

1004-
internal class TestStreamContext : ConnectionContext, IStreamDirectionFeature, IStreamIdFeature, IProtocolErrorCodeFeature, IPersistentStateFeature, IStreamAbortFeature, IConnectionCompleteFeature
1004+
internal class TestStreamContext : ConnectionContext, IStreamDirectionFeature, IStreamIdFeature, IProtocolErrorCodeFeature, IPersistentStateFeature, IStreamAbortFeature, IDisposable
10051005
{
10061006
private readonly Http3InMemory _testBase;
10071007

@@ -1017,7 +1017,6 @@ internal class TestStreamContext : ConnectionContext, IStreamDirectionFeature, I
10171017
// Persistent state collection is not reset with a stream by design.
10181018
private IDictionary<object, object> _persistentState;
10191019

1020-
private Stack<KeyValuePair<Func<object, Task>, object>> _onCompleted;
10211020
private TaskCompletionSource _disposingTcs;
10221021
private TaskCompletionSource _disposedTcs;
10231022

@@ -1067,7 +1066,6 @@ public void Initialize(long streamId)
10671066
Features.Set<IStreamAbortFeature>(this);
10681067
Features.Set<IProtocolErrorCodeFeature>(this);
10691068
Features.Set<IPersistentStateFeature>(this);
1070-
Features.Set<IConnectionCompleteFeature>(this);
10711069

10721070
StreamId = streamId;
10731071
_testBase.Logger.LogInformation($"Initializing stream {streamId}");
@@ -1134,23 +1132,25 @@ public override async ValueTask DisposeAsync()
11341132

11351133
var readerCompletedSuccessfully = _transportPipeReader.IsCompletedSuccessfully;
11361134
var writerCompletedSuccessfully = _transportPipeWriter.IsCompletedSuccessfully;
1137-
var canReuse = !_isAborted &&
1135+
CanReuse = !_isAborted &&
11381136
readerCompletedSuccessfully &&
11391137
writerCompletedSuccessfully;
11401138

11411139
_pair.Transport.Input.Complete();
11421140
_pair.Transport.Output.Complete();
1141+
}
11431142

1144-
await ConnectionCompletion.FireOnCompletedAsync(_testBase.Logger, _onCompleted);
1145-
1146-
if (canReuse)
1143+
public void Dispose()
1144+
{
1145+
if (CanReuse)
11471146
{
11481147
_testBase.Logger.LogDebug($"Pooling stream {StreamId} for reuse.");
11491148
_testBase._streamContextPool.Enqueue(this);
11501149
}
11511150
else
11521151
{
1153-
_testBase.Logger.LogDebug($"Can't reuse stream {StreamId}. Aborted: {_isAborted}, Reader completed successfully: {readerCompletedSuccessfully}, Writer completed successfully: {writerCompletedSuccessfully}.");
1152+
// Note that completed flags could be out of date at this point.
1153+
_testBase.Logger.LogDebug($"Can't reuse stream {StreamId}. Aborted: {_isAborted}, Reader completed successfully: {_transportPipeReader.IsCompletedSuccessfully}, Writer completed successfully: {_transportPipeWriter.IsCompletedSuccessfully}.");
11541154
}
11551155

11561156
Disposed = true;
@@ -1171,6 +1171,8 @@ IDictionary<object, object> IPersistentStateFeature.State
11711171
}
11721172
}
11731173

1174+
public bool CanReuse { get; private set; }
1175+
11741176
void IStreamAbortFeature.AbortRead(long errorCode, ConnectionAbortedException abortReason)
11751177
{
11761178
AbortReadException = abortReason;
@@ -1180,14 +1182,5 @@ void IStreamAbortFeature.AbortWrite(long errorCode, ConnectionAbortedException a
11801182
{
11811183
AbortWriteException = abortReason;
11821184
}
1183-
1184-
public void OnCompleted(Func<object, Task> callback, object state)
1185-
{
1186-
if (_onCompleted == null)
1187-
{
1188-
_onCompleted = new Stack<KeyValuePair<Func<object, Task>, object>>();
1189-
}
1190-
_onCompleted.Push(new KeyValuePair<Func<object, Task>, object>(callback, state));
1191-
}
11921185
}
11931186
}

‎src/Servers/Kestrel/tools/CodeGenerator/TransportConnectionFeatureCollection.cs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,7 @@ public static string GenerateFile()
2424
"IProtocolErrorCodeFeature",
2525
"IStreamDirectionFeature",
2626
"IStreamIdFeature",
27-
"IStreamAbortFeature",
28-
"IConnectionCompleteFeature"
27+
"IStreamAbortFeature"
2928
};
3029

3130
var implementedFeatures = new[]

0 commit comments

Comments
 (0)