Add PersistentState

Author: JamesNK

src/Servers/Connections.Abstractions/src/ConnectionContext.cs

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Generic;
 using System.IO.Pipelines;
 using Microsoft.AspNetCore.Connections.Features;
 
@@ -12,6 +13,30 @@ namespace Microsoft.AspNetCore.Connections
     /// </summary>
     public abstract class ConnectionContext : BaseConnectionContext, IAsyncDisposable
     {
+        internal IDictionary<object, object?>? _persistentState;
+
+        /// <summary>
+        /// Gets or sets a key/value collection that can be used to persist state between connections.
+        /// Whether a transport pools and reuses <see cref="ConnectionContext"/> instances and allows state to
+        /// be persisted depends on the transport implementation.
+        /// <para>
+        /// Because values added to persistent state can live in memory until a <see cref="ConnectionContext"/>
+        /// is no longer pooled, use caution with this collection to avoid excessive memory use.
+        /// </para>
+        /// </summary>
+        public virtual IDictionary<object, object?> PersistentState
+        {
+            get
+            {
+                // Lazily allocate connection metadata
+                return _persistentState ?? (_persistentState = new ConnectionItems());
+            }
+            set
+            {
+                _persistentState = value;
+            }
+        }
+
         /// <summary>
         /// Gets or sets the <see cref="IDuplexPipe"/> that can be used to read or write data on this connection.
         /// </summary>

src/Servers/Connections.Abstractions/src/PublicAPI.Unshipped.txt

@@ -25,3 +25,5 @@ Microsoft.AspNetCore.Connections.MultiplexedConnectionContext.MultiplexedConnect
 Microsoft.AspNetCore.Connections.MultiplexedConnectionDelegate
 abstract Microsoft.AspNetCore.Connections.MultiplexedConnectionContext.AcceptAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask<Microsoft.AspNetCore.Connections.ConnectionContext?>
 abstract Microsoft.AspNetCore.Connections.MultiplexedConnectionContext.ConnectAsync(Microsoft.AspNetCore.Http.Features.IFeatureCollection? features = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask<Microsoft.AspNetCore.Connections.ConnectionContext!>
+virtual Microsoft.AspNetCore.Connections.ConnectionContext.PersistentState.get -> System.Collections.Generic.IDictionary<object!, object?>!
+virtual Microsoft.AspNetCore.Connections.ConnectionContext.PersistentState.set -> void

src/Servers/Kestrel/Core/src/Internal/Http3/Http3Connection.cs

@@ -23,6 +23,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http3
 {
     internal class Http3Connection : IHttp3StreamLifetimeHandler, IRequestProcessor
     {
+        private static readonly object StreamPersistentStateKey = new object();
+
         // Internal for unit testing
         internal readonly Dictionary<long, IHttp3Stream> _streams = new Dictionary<long, IHttp3Stream>();
         internal IHttp3StreamLifetimeHandler _streamLifetimeHandler;
@@ -270,16 +272,18 @@ public async Task ProcessRequestsAsync<TContext>(IHttpApplication<TContext> appl
                             // Request stream
                             UpdateHighestStreamId(streamIdFeature.StreamId);
 
+                            Http3Stream<TContext> stream;
+
                             // Check whether there is an existing HTTP/3 stream on the transport stream.
                             // A stream will only be cached if the transport stream itself is reused.
-                            var stream = streamContext.Features.Get<ICachedHttp3StreamFeature<TContext>>()?.CachedStream;
-                            if (stream == null)
+                            if (!streamContext.PersistentState.TryGetValue(StreamPersistentStateKey, out var s))
                             {
                                 stream = new Http3Stream<TContext>(application, CreateHttpStreamContext(streamContext));
-                                streamContext.Features.Set<ICachedHttp3StreamFeature<TContext>>(new DefaultCachedHttp3StreamFeature<TContext>(stream));
+                                streamContext.PersistentState.Add(StreamPersistentStateKey, stream);
                             }
                             else
                             {
+                                stream = (Http3Stream<TContext>)s!;
                                 stream.InitializeWithExistingContext(streamContext.Transport);
                             }
 

src/Servers/Kestrel/Transport.Quic/src/Internal/QuicConnectionContext.cs

@@ -94,6 +94,11 @@ public override void Abort(ConnectionAbortedException abortReason)
                 {
                     context = new QuicStreamContext(this, _context);
                 }
+                else
+                {
+                    context.ResetFeatureCollection();
+                    context.ResetItems();
+                }
 
                 context.Initialize(stream);
                 context.Start();

src/Servers/Kestrel/shared/TransportConnection.cs

@@ -54,6 +54,11 @@ public override string ConnectionId
             }
         }
 
+        internal void ResetItems()
+        {
+            _items?.Clear();
+        }
+
         public override CancellationToken ConnectionClosed { get; set; }
 
         // DO NOT remove this override to ConnectionContext.Abort. Doing so would cause