React to aspnet/HttpAbstractions#160 - Implemented OnResponseCompleted

Author: ajaybhargavb

src/Microsoft.AspNet.Server.WebListener/FeatureContext.cs

@@ -360,6 +360,11 @@ void IHttpResponseFeature.OnSendingHeaders(Action<object> callback, object state
             Response.OnSendingHeaders(callback, state);
         }
 
+        void IHttpResponseFeature.OnResponseCompleted(Action<object> callback, object state)
+        {
+            Response.OnResponseCompleted(callback, state);
+        }
+
         string IHttpResponseFeature.ReasonPhrase
         {
             get { return Response.ReasonPhrase; }

src/Microsoft.Net.Http.Server/RequestProcessing/Response.cs

@@ -48,6 +48,7 @@ public sealed unsafe class Response
         private BoundaryType _boundaryType;
         private UnsafeNclNativeMethods.HttpApi.HTTP_RESPONSE_V2 _nativeResponse;
         private IList<Tuple<Action<object>, object>> _onSendingHeadersActions;
+        private IList<Tuple<Action<object>, object>> _onResponseCompletedActions;
 
         private RequestContext _requestContext;
 
@@ -63,6 +64,7 @@ internal Response(RequestContext httpContext)
             _nativeResponse.Response_V1.Version.MinorVersion = 1;
             _responseState = ResponseState.Created;
             _onSendingHeadersActions = new List<Tuple<Action<object>, object>>();
+            _onResponseCompletedActions = new List<Tuple<Action<object>, object>>();
         }
 
         private enum ResponseState
@@ -284,6 +286,7 @@ internal void Dispose()
             {
                 return;
             }
+            NotifyOnResponseCompleted();
             // TODO: Verbose log
             EnsureResponseStream();
             _nativeStream.Dispose();
@@ -829,6 +832,17 @@ public void OnSendingHeaders(Action<object> callback, object state)
             actions.Add(new Tuple<Action<object>, object>(callback, state));
         }
 
+        public void OnResponseCompleted(Action<object> callback, object state)
+        {
+            var actions = _onResponseCompletedActions;
+            if (actions == null)
+            {
+                throw new InvalidOperationException("Response already completed");
+            }
+
+            actions.Add(new Tuple<Action<object>, object>(callback, state));
+        }
+
         private void NotifyOnSendingHeaders()
         {
             var actions = Interlocked.Exchange(ref _onSendingHeadersActions, null);
@@ -845,6 +859,30 @@ private void NotifyOnSendingHeaders()
             }
         }
 
+        private void NotifyOnResponseCompleted()
+        {
+            var actions = Interlocked.Exchange(ref _onResponseCompletedActions, null);
+            if (actions == null)
+            {
+                // Something threw the first time, do not try again.
+                return;
+            }
+
+            foreach (var actionPair in actions)
+            {
+                try
+                {
+                    actionPair.Item1(actionPair.Item2);
+                }
+                catch (Exception ex)
+                {
+                    RequestContext.Logger.LogWarning(
+                        String.Format(Resources.Warning_ExceptionInOnResponseCompletedAction, nameof(OnResponseCompleted)),
+                        ex);
+                }
+            }
+        }
+
         private class SendResponseLogContext : ReflectionBasedLogValues
         {
             private readonly Response _response;

src/Microsoft.Net.Http.Server/Resources.Designer.cs

@@ -147,4 +147,7 @@
   <data name="Exception_WrongIAsyncResult" xml:space="preserve">
     <value>The given IAsyncResult does not match this opperation.</value>
   </data>
+  <data name="Warning_ExceptionInOnResponseCompletedAction" xml:space="preserve">
+    <value>An exception occured while running an action registered with {0}.</value>
+  </data>
 </root>