React to aspnet/HttpAbstractions#160 - Implementing OnResponseCompleted

Author: ajaybhargavb

src/Kestrel/ServerRequest.cs

@@ -204,6 +204,11 @@ void IHttpResponseFeature.OnSendingHeaders(Action<object> callback, object state
             _frame.OnSendingHeaders(callback, state);
         }
 
+        void IHttpResponseFeature.OnResponseCompleted(Action<object> callback, object state)
+        {
+            _frame.OnResponseCompleted(callback, state);
+        }
+
         bool IHttpUpgradeFeature.IsUpgradableRequest
         {
             get

src/Microsoft.AspNet.Server.Kestrel/Http/Frame.cs

@@ -64,7 +64,9 @@ enum Mode
         */
 
         List<KeyValuePair<Action<object>, object>> _onSendingHeaders;
+        List<KeyValuePair<Action<object>, object>> _onResponseCompleted;
         object _onSendingHeadersSync = new Object();
+        object _onResponseCompletedSync = new Object();
 
         public Frame(ConnectionContext context) : base(context)
         {
@@ -200,6 +202,18 @@ public void OnSendingHeaders(Action<object> callback, object state)
             }
         }
 
+        public void OnResponseCompleted(Action<object> callback, object state)
+        {
+            lock (_onResponseCompletedSync)
+            {
+                if (_onResponseCompleted == null)
+                {
+                    _onResponseCompleted = new List<KeyValuePair<Action<object>, object>>();
+                }
+                _onResponseCompleted.Add(new KeyValuePair<Action<object>, object>(callback, state));
+            }
+        }
+
         private void FireOnSendingHeaders()
         {
             List<KeyValuePair<Action<object>, object>> onSendingHeaders = null;
@@ -217,6 +231,30 @@ private void FireOnSendingHeaders()
             }
         }
 
+        private void FireOnResponseCompleted()
+        {
+            List<KeyValuePair<Action<object>, object>> onResponseCompleted = null;
+            lock (_onResponseCompletedSync)
+            {
+                onResponseCompleted = _onResponseCompleted;
+                _onResponseCompleted = null;
+            }
+            if (onResponseCompleted != null)
+            {
+                foreach (var entry in onResponseCompleted)
+                {
+                    try
+                    {
+                        entry.Key.Invoke(entry.Value);
+                    }
+                    catch
+                    {
+                        // Ignore exceptions
+                    }
+                }
+            }
+        }
+
         private async Task ExecuteAsync()
         {
             Exception error = null;
@@ -230,6 +268,7 @@ private async Task ExecuteAsync()
             }
             finally
             {
+                FireOnResponseCompleted();
                 ProduceEnd(error);
             }
         }