General cleanup for client disconnect handling. (#1446)

src/AspNetCoreModuleV2/AspNetCore/DisconnectHandler.cpp

@@ -10,15 +10,16 @@ void DisconnectHandler::NotifyDisconnect()
 {
     try
     {
-        std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> module;
+        std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> pHandler;
         {
             SRWExclusiveLock lock(m_handlerLock);
-            m_pHandler.swap(module);
+            m_pHandler.swap(pHandler);
+            m_disconnectFired = true;
         }
 
-        if (module != nullptr)
+        if (pHandler != nullptr)
         {
-            module->NotifyDisconnect();
+            pHandler->NotifyDisconnect();
         }
     }
     catch (...)
@@ -29,12 +30,28 @@ void DisconnectHandler::NotifyDisconnect()
 
 void DisconnectHandler::CleanupStoredContext() noexcept
 {
-    SetHandler(nullptr);
     delete this;
 }
 
-void DisconnectHandler::SetHandler(std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> handler) noexcept
+void DisconnectHandler::SetHandler(std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> handler)
 {
-    SRWExclusiveLock lock(m_handlerLock);
+    IREQUEST_HANDLER* pHandler = nullptr;
-    handler.swap(m_pHandler);
+    {
+        SRWExclusiveLock lock(m_handlerLock);
+
+        handler.swap(m_pHandler);
+        pHandler = m_pHandler.get();
+    }
+
+    assert(pHandler != nullptr);
+
+    if (pHandler != nullptr && (m_disconnectFired || m_pHttpConnection != nullptr && !m_pHttpConnection->IsConnected()))
+    {
+        pHandler->NotifyDisconnect();
+    }
+}
+
+void DisconnectHandler::RemoveHandler() noexcept
+{
+    m_pHandler = nullptr;
 }

src/AspNetCoreModuleV2/AspNetCore/DisconnectHandler.h

@@ -11,8 +11,8 @@ class ASPNET_CORE_PROXY_MODULE;
 class DisconnectHandler final: public IHttpConnectionStoredContext
 {
 public:
-    DisconnectHandler()
+    DisconnectHandler(IHttpConnection* pHttpConnection)
-        : m_pHandler(nullptr)
+        : m_pHandler(nullptr), m_pHttpConnection(pHttpConnection), m_disconnectFired(false)
     {
         InitializeSRWLock(&m_handlerLock);
     }
@@ -20,7 +20,7 @@ public:
     virtual
     ~DisconnectHandler()
     {
-        SetHandler(nullptr);
+        RemoveHandler();
     }
 
     void
@@ -30,10 +30,14 @@ public:
     CleanupStoredContext() noexcept override;
 
     void
-    SetHandler(std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> handler) noexcept;
+    SetHandler(std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> handler);
+
+    void RemoveHandler() noexcept;
 
 private:
     SRWLOCK m_handlerLock {};
     std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> m_pHandler;
+    IHttpConnection* m_pHttpConnection;
+    bool m_disconnectFired;
 };
 

src/AspNetCoreModuleV2/AspNetCore/proxymodule.cpp

@@ -96,7 +96,6 @@ ASPNET_CORE_PROXY_MODULE::OnExecuteRequestHandler(
             FINISHED(HRESULT_FROM_WIN32(ERROR_SERVER_SHUTDOWN_IN_PROGRESS));
         }
 
-
         FINISHED_IF_FAILED(m_pApplicationManager->GetOrCreateApplicationInfo(
             *pHttpContext,
             m_pApplicationInfo));
@@ -185,7 +184,7 @@ void ASPNET_CORE_PROXY_MODULE::SetupDisconnectHandler(IHttpContext * pHttpContex
 
     if (pDisconnectHandler == nullptr)
     {
-        auto newHandler = std::make_unique<DisconnectHandler>();
+        auto newHandler = std::make_unique<DisconnectHandler>(pHttpContext->GetConnection());
         pDisconnectHandler = newHandler.get();
         // ModuleContextContainer takes ownership of disconnectHandler
         // we are trusting that it would not release it before deleting the context
@@ -207,6 +206,6 @@ void ASPNET_CORE_PROXY_MODULE::RemoveDisconnectHandler() noexcept
 
     if (handler != nullptr)
     {
-        handler->SetHandler(nullptr);
+        handler->RemoveHandler();
     }
 }

src/AspNetCoreModuleV2/InProcessRequestHandler/inprocesshandler.cpp

@@ -26,8 +26,10 @@ IN_PROCESS_HANDLER::IN_PROCESS_HANDLER(
    m_pRequestHandlerContext(pRequestHandlerContext),
    m_pAsyncCompletionHandler(pAsyncCompletion),
    m_pDisconnectHandler(pDisconnectHandler),
-   m_moduleTracer(pW3Context->GetTraceContext())
+   m_moduleTracer(pW3Context->GetTraceContext()),
+   m_disconnectFired(false)
 {
+    InitializeSRWLock(&m_srwDisconnectLock);
 }
 
 __override
@@ -92,16 +94,29 @@ REQUEST_NOTIFICATION_STATUS IN_PROCESS_HANDLER::ServerShutdownMessage()
 VOID
 IN_PROCESS_HANDLER::NotifyDisconnect()
 {
+    m_moduleTracer.RequestDisconnect();
+
     if (m_pApplication->QueryBlockCallbacksIntoManaged() ||
         m_fManagedRequestComplete)
     {
         return;
     }
 
-    m_moduleTracer.RequestDisconnect();
+    // NotifyDisconnect can be called before the m_pManagedHttpContext is set,
+    // so save that in a bool.
+    // Don't lock when calling m_pDisconnect to avoid the potential deadlock between this
+    // and SetManagedHttpContext
+    void* pManagedHttpContext = nullptr;
+    {
+        SRWExclusiveLock lock(m_srwDisconnectLock);
+        pManagedHttpContext = m_pManagedHttpContext;
+        m_disconnectFired = true;
+    }
 
-    assert(m_pManagedHttpContext != nullptr);
+    if (pManagedHttpContext != nullptr)
-    m_pDisconnectHandler(m_pManagedHttpContext);
+    {
+        m_pDisconnectHandler(m_pManagedHttpContext);
+    }
 }
 
 VOID
@@ -127,7 +142,15 @@ IN_PROCESS_HANDLER::SetManagedHttpContext(
     PVOID pManagedHttpContext
 )
 {
-    m_pManagedHttpContext = pManagedHttpContext;
+    {
+        SRWExclusiveLock lock(m_srwDisconnectLock);
+        m_pManagedHttpContext = pManagedHttpContext;
+    }
+
+    if (m_disconnectFired && m_pManagedHttpContext != nullptr)
+    {
+        m_pDisconnectHandler(m_pManagedHttpContext);
+    }
 }
 
 // static

src/AspNetCoreModuleV2/InProcessRequestHandler/inprocesshandler.h

@@ -85,4 +85,6 @@ private:
     PFN_DISCONNECT_HANDLER       m_pDisconnectHandler;
     static ALLOC_CACHE_HANDLER *   sm_pAlloc;
     ModuleTracer m_moduleTracer;
+    bool m_disconnectFired;
+    SRWLOCK m_srwDisconnectLock;
 };

src/Microsoft.AspNetCore.Server.IIS/Core/IISHttpContext.IO.cs

@@ -190,13 +190,20 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
 
             _bodyOutput.Dispose();
 
-            try
+            var cts = _abortedCts;
+            if (cts != null)
             {
-                _abortedCts?.Cancel();
+                ThreadPool.QueueUserWorkItem(t =>
-            }
+                {
-            catch (Exception)
+                    try
-            {
+                    {
-                // ignore
+                        cts.Cancel();
+                    }
+                    catch (Exception)
+                    {
+                        // ignore
+                    }
+                });
             }
         }
 

src/Microsoft.AspNetCore.Server.IIS/Core/IISHttpContext.cs

@@ -109,8 +109,6 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
         {
             _thisHandle = GCHandle.Alloc(this);
 
-            NativeMethods.HttpSetManagedContext(_pInProcessHandler, (IntPtr)_thisHandle);
-
             Method = GetVerb();
 
             RawTarget = GetRawUrl();
@@ -170,6 +168,8 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
                     resumeWriterThreshold: ResumeWriterTheshold,
                     minimumSegmentSize: MinAllocBufferSize));
             _bodyOutput = new OutputProducer(pipe);
+
+            NativeMethods.HttpSetManagedContext(_pInProcessHandler, (IntPtr)_thisHandle);
         }
 
         private string GetOriginalPath()

src/Microsoft.AspNetCore.Server.IIS/Core/IISHttpServer.cs

@@ -213,7 +213,6 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
             {
                 context?.Server._logger.LogError(0, ex, $"Unexpected exception in {nameof(IISHttpServer)}.{nameof(OnDisconnect)}.");
             }
-
         }
 
         private static NativeMethods.REQUEST_NOTIFICATION_STATUS OnAsyncCompletion(IntPtr pvManagedHttpContext, int hr, int bytes)

test/Common.FunctionalTests/Inprocess/ClientDisconnectTests.cs

@@ -1,8 +1,12 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // 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.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Testing.xunit;
+using Microsoft.Extensions.Logging;
 using Xunit;
 
 namespace Microsoft.AspNetCore.Server.IISIntegration.FunctionalTests
@@ -61,5 +65,44 @@ namespace Microsoft.AspNetCore.Server.IISIntegration.FunctionalTests
 
             await _fixture.Client.RetryRequestAsync("/WaitingRequestCount", async message => await message.Content.ReadAsStringAsync() == "0");
         }
+
+        [ConditionalFact]
+        public async Task ClientDisconnectCallbackStress()
+        {
+            // Fixture initialization fails if inside of the Task.Run, so send an
+            // initial request to initialize the fixture.
+            var response = await _fixture.Client.GetAsync("HelloWorld");
+            var numTotalRequests = 0;
+            for (var j = 0; j < 20; j++)
+            {
+                // Windows has a max connection limit of 10 for the IIS server,
+                // so setting limit fairly low. 
+                const int numRequests = 5;
+                async Task RunRequests()
+                {
+                    using (var connection = _fixture.CreateTestConnection())
+                    {
+                        await connection.Send(
+                            "GET /WaitForAbort HTTP/1.1",
+                            "Host: localhost",
+                            "Connection: close",
+                            "",
+                            "");
+                        await _fixture.Client.RetryRequestAsync("/WaitingRequestCount", async message => await message.Content.ReadAsStringAsync() != "0");
+                        Interlocked.Increment(ref numTotalRequests);
+                    }
+                }
+
+                List<Task> tasks = new List<Task>();
+                for (int i = 0; i < numRequests; i++)
+                {
+                    tasks.Add(Task.Run(RunRequests));
+                }
+
+                await Task.WhenAll(tasks);
+
+                await _fixture.Client.RetryRequestAsync("/WaitingRequestCount", async message => await message.Content.ReadAsStringAsync() == "0");
+            }
+        }
     }
 }