TransportConnection IFeatureCollection

aspnet · Feb 10, 2018 · 5b345fa · 5b345fa
1 parent f8f95ef
commit 5b345fa
Show file tree

Hide file tree

Showing 3 changed files with 77 additions and 47 deletions.
diff --git a/src/Kestrel.Core/Internal/Http/HttpProtocol.Generated.cs b/src/Kestrel.Core/Internal/Http/HttpProtocol.Generated.cs
@@ -453,17 +453,16 @@ TFeature IFeatureCollection.Get<TFeature>()
             {
                 feature = (TFeature)_currentIHttpSendFileFeature;
             }
-            else
+            else if (MaybeExtra != null)
             {
                 feature = (TFeature)(ExtraFeatureGet(typeof(TFeature)));
             }
-
-            if (feature != null)
+            else
             {
-                return feature;
+                feature = ConnectionFeatures.Get<TFeature>();
             }
 
-            return (TFeature)ConnectionFeatures[typeof(TFeature)];
+            return feature;
         }
 
         private IEnumerable<KeyValuePair<Type, object>> FastEnumerable()

diff --git a/src/Kestrel.Transport.Abstractions/Internal/TransportConnection.Features.cs b/src/Kestrel.Transport.Abstractions/Internal/TransportConnection.Features.cs
@@ -112,69 +112,101 @@ IDuplexPipe IConnectionTransportFeature.Application
 
         object IFeatureCollection.this[Type key]
         {
-            get => FastFeatureGet(key);
-            set => FastFeatureSet(key, value);
-        }
+            get
+            {
+                if (key == IHttpConnectionFeatureType)
+                {
+                    return _currentIHttpConnectionFeature;
+                }
 
-        TFeature IFeatureCollection.Get<TFeature>()
-        {
-            return (TFeature)FastFeatureGet(typeof(TFeature));
-        }
+                if (key == IConnectionIdFeatureType)
+                {
+                    return _currentIConnectionIdFeature;
+                }
 
-        void IFeatureCollection.Set<TFeature>(TFeature instance)
-        {
-            FastFeatureSet(typeof(TFeature), instance);
-        }
+                if (key == IConnectionTransportFeatureType)
+                {
+                    return _currentIConnectionTransportFeature;
+                }
 
-        IEnumerator<KeyValuePair<Type, object>> IEnumerable<KeyValuePair<Type, object>>.GetEnumerator() => FastEnumerable().GetEnumerator();
+                if (MaybeExtra != null)
+                {
+                    return ExtraFeatureGet(key);
+                }
 
-        IEnumerator IEnumerable.GetEnumerator() => FastEnumerable().GetEnumerator();
+                return null;
+            }
+            set
+            {
+                _featureRevision++;
 
-        private object FastFeatureGet(Type key)
+                if (key == IHttpConnectionFeatureType)
+                {
+                    _currentIHttpConnectionFeature = value;
+                }
+                else if (key == IConnectionIdFeatureType)
+                {
+                    _currentIConnectionIdFeature = value;
+                }
+                else if (key == IConnectionTransportFeatureType)
+                {
+                    _currentIConnectionTransportFeature = value;
+                }
+                else
+                {
+                    ExtraFeatureSet(key, value);
+                }
+            }
+        }
+
+        TFeature IFeatureCollection.Get<TFeature>()
         {
-            if (key == IHttpConnectionFeatureType)
+            if (typeof(TFeature) == typeof(IHttpConnectionFeature))
             {
-                return _currentIHttpConnectionFeature;
+                return (TFeature)_currentIHttpConnectionFeature;
             }
-
-            if (key == IConnectionIdFeatureType)
+            else if (typeof(TFeature) == typeof(IConnectionIdFeature))
             {
-                return _currentIConnectionIdFeature;
+                return (TFeature)_currentIConnectionIdFeature;
             }
-
-            if (key == IConnectionTransportFeatureType)
+            else if (typeof(TFeature) == typeof(IConnectionTransportFeature))
             {
-                return _currentIConnectionTransportFeature;
+                return (TFeature)_currentIConnectionTransportFeature;
+            }
+            else if (MaybeExtra != null)
+            {
+                return (TFeature)ExtraFeatureGet(typeof(TFeature));
             }
 
-            return ExtraFeatureGet(key);
+            return default;
         }
 
-        private void FastFeatureSet(Type key, object feature)
+        void IFeatureCollection.Set<TFeature>(TFeature instance)
         {
             _featureRevision++;
 
-            if (key == IHttpConnectionFeatureType)
+            if (typeof(TFeature) == typeof(IHttpConnectionFeature))
             {
-                _currentIHttpConnectionFeature = feature;
-                return;
+                _currentIHttpConnectionFeature = instance;
             }
-
-            if (key == IConnectionIdFeatureType)
+            else if (typeof(TFeature) == typeof(IConnectionIdFeature))
             {
-                _currentIConnectionIdFeature = feature;
-                return;
+                _currentIConnectionIdFeature = instance;
             }
-
-            if (key == IConnectionTransportFeatureType)
+            else if (typeof(TFeature) == typeof(IConnectionTransportFeature))
             {
-                _currentIConnectionTransportFeature = feature;
-                return;
+                _currentIConnectionTransportFeature = instance;
+            }
+            else
+            {
+                ExtraFeatureSet(typeof(TFeature), instance);
             }
-
-            ExtraFeatureSet(key, feature);
         }
 
+        IEnumerator<KeyValuePair<Type, object>> IEnumerable<KeyValuePair<Type, object>>.GetEnumerator() => FastEnumerable().GetEnumerator();
+
+        IEnumerator IEnumerable.GetEnumerator() => FastEnumerable().GetEnumerator();
+
         private IEnumerable<KeyValuePair<Type, object>> FastEnumerable()
         {
             if (_currentIHttpConnectionFeature != null)

diff --git a/tools/CodeGenerator/HttpProtocolFeatureCollection.cs b/tools/CodeGenerator/HttpProtocolFeatureCollection.cs
@@ -165,17 +165,16 @@ TFeature IFeatureCollection.Get<TFeature>()
             {{
                 feature = (TFeature)_current{feature.Name};
             }}")}
-            else
+            else if (MaybeExtra != null)
             {{
                 feature = (TFeature)(ExtraFeatureGet(typeof(TFeature)));
             }}
-
-            if (feature != null)
+            else
             {{
-                return feature;
+                feature = ConnectionFeatures.Get<TFeature>();
             }}
 
-            return (TFeature)ConnectionFeatures[typeof(TFeature)];
+            return feature;
         }}
 
         private IEnumerable<KeyValuePair<Type, object>> FastEnumerable()