Disconnect the disconnect handler when request processing ends (#1471)

src/AspNetCoreModuleV2/AspNetCore/DisconnectHandler.cpp

@@ -4,15 +4,21 @@
 #include "DisconnectHandler.h"
 #include "exceptions.h"
 #include "proxymodule.h"
+#include "SRWExclusiveLock.h"
 
 void DisconnectHandler::NotifyDisconnect()
 {
     try
     {
-        const auto module = m_pModule.exchange(nullptr);
+        std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> module;
+        {
+            SRWExclusiveLock lock(m_handlerLock);
+            m_pHandler.swap(module);
+        }
+
         if (module != nullptr)
         {
-            module ->NotifyDisconnect();
+            module->NotifyDisconnect();
         }
     }
     catch (...)
@@ -27,7 +33,8 @@ void DisconnectHandler::CleanupStoredContext() noexcept
     delete this;
 }
 
-void DisconnectHandler::SetHandler(ASPNET_CORE_PROXY_MODULE * module) noexcept
+void DisconnectHandler::SetHandler(std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> handler) noexcept
 {
-    m_pModule = module;
+    SRWExclusiveLock lock(m_handlerLock);
+    handler.swap(m_pHandler);
 }

src/AspNetCoreModuleV2/AspNetCore/DisconnectHandler.h

@@ -2,7 +2,9 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 #pragma once
-#include <atomic>
+
+#include <memory>
+#include "irequesthandler.h"
 
 class ASPNET_CORE_PROXY_MODULE;
 
@@ -10,8 +12,9 @@ class DisconnectHandler final: public IHttpConnectionStoredContext
 {
 public:
     DisconnectHandler()
-        : m_pModule(nullptr)
+        : m_pHandler(nullptr)
     {
+        InitializeSRWLock(&m_handlerLock);
     }
 
     virtual
@@ -27,9 +30,10 @@ public:
     CleanupStoredContext() noexcept override;
 
     void
-    SetHandler(ASPNET_CORE_PROXY_MODULE * module) noexcept;
+    SetHandler(std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> handler) noexcept;
 
 private:
-    std::atomic<ASPNET_CORE_PROXY_MODULE*> m_pModule;
+    SRWLOCK m_handlerLock {};
+    std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> m_pHandler;
 };
 

src/AspNetCoreModuleV2/AspNetCore/proxymodule.cpp

@@ -71,10 +71,7 @@ ASPNET_CORE_PROXY_MODULE::ASPNET_CORE_PROXY_MODULE(HTTP_MODULE_ID moduleId, std:
 
 ASPNET_CORE_PROXY_MODULE::~ASPNET_CORE_PROXY_MODULE()
 {
-    if (m_pDisconnectHandler != nullptr)
+    RemoveDisconnectHandler();
-    {
-        m_pDisconnectHandler->SetHandler(nullptr);
-    }
 }
 
 __override
@@ -94,25 +91,6 @@ ASPNET_CORE_PROXY_MODULE::OnExecuteRequestHandler(
             FINISHED(HRESULT_FROM_WIN32(ERROR_SERVER_SHUTDOWN_IN_PROGRESS));
         }
 
-        auto moduleContainer = pHttpContext
-            ->GetConnection()
-            ->GetModuleContextContainer();
-
-        #pragma warning( push )
-        #pragma warning ( disable : 26466 ) // Disable "Don't use static_cast downcasts". We build without RTTI support so dynamic_cast is not available
-        m_pDisconnectHandler = static_cast<DisconnectHandler*>(moduleContainer->GetConnectionModuleContext(m_moduleId));
-        #pragma warning( push )
-
-        if (m_pDisconnectHandler == nullptr)
-        {
-            auto disconnectHandler = std::make_unique<DisconnectHandler>();
-            m_pDisconnectHandler = disconnectHandler.get();
-            // ModuleContextContainer takes ownership of disconnectHandler
-            // we are trusting that it would not release it before deleting the context
-            FINISHED_IF_FAILED(moduleContainer->SetConnectionModuleContext(static_cast<IHttpConnectionStoredContext*>(disconnectHandler.release()), m_moduleId));
-        }
-
-        m_pDisconnectHandler->SetHandler(this);
 
         FINISHED_IF_FAILED(m_pApplicationManager->GetOrCreateApplicationInfo(
             *pHttpContext,
@@ -120,6 +98,8 @@ ASPNET_CORE_PROXY_MODULE::OnExecuteRequestHandler(
 
         FINISHED_IF_FAILED(m_pApplicationInfo->CreateHandler(*pHttpContext, m_pHandler));
 
+        SetupDisconnectHandler(pHttpContext);
+
         retVal = m_pHandler->OnExecuteRequestHandler();
     }
     catch (...)
@@ -141,7 +121,7 @@ Finished:
         }
     }
 
-    return retVal;
+    return HandleNotificationStatus(retVal);
 }
 
 __override
@@ -156,18 +136,62 @@ ASPNET_CORE_PROXY_MODULE::OnAsyncCompletion(
 {
     try
     {
-        return m_pHandler->OnAsyncCompletion(
+        return HandleNotificationStatus(m_pHandler->OnAsyncCompletion(
             pCompletionInfo->GetCompletionBytes(),
-            pCompletionInfo->GetCompletionStatus());
+            pCompletionInfo->GetCompletionStatus()));
     }
     catch (...)
     {
         OBSERVE_CAUGHT_EXCEPTION();
-        return RQ_NOTIFICATION_FINISH_REQUEST;
+        return HandleNotificationStatus(RQ_NOTIFICATION_FINISH_REQUEST);
     }
 }
 
-void ASPNET_CORE_PROXY_MODULE::NotifyDisconnect() const
+REQUEST_NOTIFICATION_STATUS ASPNET_CORE_PROXY_MODULE::HandleNotificationStatus(REQUEST_NOTIFICATION_STATUS status) noexcept
 {
-    m_pHandler->NotifyDisconnect();
+    if (status != RQ_NOTIFICATION_PENDING)
+    {
+        RemoveDisconnectHandler();
+    }
+
+    return status;
+}
+
+void ASPNET_CORE_PROXY_MODULE::SetupDisconnectHandler(IHttpContext * pHttpContext)
+{
+    auto moduleContainer = pHttpContext
+        ->GetConnection()
+        ->GetModuleContextContainer();
+
+    #pragma warning( push )
+    #pragma warning ( disable : 26466 ) // Disable "Don't use static_cast downcasts". We build without RTTI support so dynamic_cast is not available
+    auto pDisconnectHandler = static_cast<DisconnectHandler*>(moduleContainer->GetConnectionModuleContext(m_moduleId));
+    #pragma warning( push )
+
+    if (pDisconnectHandler == nullptr)
+    {
+        auto newHandler = std::make_unique<DisconnectHandler>();
+        pDisconnectHandler = newHandler.get();
+        // ModuleContextContainer takes ownership of disconnectHandler
+        // we are trusting that it would not release it before deleting the context
+        LOG_IF_FAILED(moduleContainer->SetConnectionModuleContext(static_cast<IHttpConnectionStoredContext*>(newHandler.release()), m_moduleId));
+    }
+
+    // make code analysis happy
+    if (pDisconnectHandler != nullptr)
+    {
+        pDisconnectHandler->SetHandler(::ReferenceRequestHandler(m_pHandler.get()));
+        m_pDisconnectHandler = pDisconnectHandler;
+    }
+}
+
+void ASPNET_CORE_PROXY_MODULE::RemoveDisconnectHandler() noexcept
+{
+    auto handler = m_pDisconnectHandler;
+    m_pDisconnectHandler = nullptr;
+
+    if (handler != nullptr)
+    {
+        handler->SetHandler(nullptr);
+    }
 }

src/AspNetCoreModuleV2/AspNetCore/proxymodule.h

@@ -47,10 +47,14 @@ class ASPNET_CORE_PROXY_MODULE : NonCopyable, public CHttpModule
         IHttpCompletionInfo *   pCompletionInfo
     ) override;
 
-    void
-    NotifyDisconnect() const;
 
  private:
+    REQUEST_NOTIFICATION_STATUS
+    HandleNotificationStatus(REQUEST_NOTIFICATION_STATUS status) noexcept;
+
+    void SetupDisconnectHandler(IHttpContext * pHttpContext);
+    void RemoveDisconnectHandler() noexcept;
+
     std::shared_ptr<APPLICATION_MANAGER> m_pApplicationManager;
     std::shared_ptr<APPLICATION_INFO> m_pApplicationInfo;
     std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> m_pHandler;

src/AspNetCoreModuleV2/CommonLib/SRWExclusiveLock.cpp

@@ -3,7 +3,7 @@
 
 #include "SRWExclusiveLock.h"
 
-SRWExclusiveLock::SRWExclusiveLock(const SRWLOCK& lock)
+SRWExclusiveLock::SRWExclusiveLock(const SRWLOCK& lock) noexcept
     : m_lock(lock)
 {
     AcquireSRWLockExclusive(const_cast<SRWLOCK*>(&m_lock));

src/AspNetCoreModuleV2/CommonLib/SRWExclusiveLock.h

@@ -8,7 +8,7 @@
 class SRWExclusiveLock
 {
 public:
-	SRWExclusiveLock(const SRWLOCK& lock);
+	SRWExclusiveLock(const SRWLOCK& lock) noexcept;
 	~SRWExclusiveLock();
 private:
     const SRWLOCK& m_lock;

src/AspNetCoreModuleV2/CommonLib/irequesthandler.h

@@ -54,3 +54,9 @@ struct IREQUEST_HANDLER_DELETER
         application->DereferenceRequestHandler();
     }
 };
+
+inline std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER> ReferenceRequestHandler(IREQUEST_HANDLER* handler)
+{
+    handler->ReferenceRequestHandler();
+    return std::unique_ptr<IREQUEST_HANDLER, IREQUEST_HANDLER_DELETER>(handler);
+};

test/Common.FunctionalTests/ClientDisconnectStress.cs

@@ -0,0 +1,65 @@
+// 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.Net.Http;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Server.IntegrationTesting;
+using Microsoft.AspNetCore.Testing.xunit;
+using Xunit;
+
+namespace Microsoft.AspNetCore.Server.IISIntegration.FunctionalTests
+{
+    [Collection(PublishedSitesCollection.Name)]
+    public class ClientDisconnectStressTests: FunctionalTestsBase
+    {
+        private readonly PublishedSitesFixture _fixture;
+
+        public ClientDisconnectStressTests(PublishedSitesFixture fixture)
+        {
+            _fixture = fixture;
+        }
+
+        [ConditionalTheory]
+        [InlineData(HostingModel.InProcess)]
+        [InlineData(HostingModel.OutOfProcess)]
+        public async Task ClientDisconnectStress(HostingModel hostingModel)
+        {
+            var site = await StartAsync(_fixture.GetBaseDeploymentParameters(hostingModel));
+            var maxRequestSize = 1000;
+            var blockSize = 40;
+            var random = new Random();
+            async Task RunRequests()
+            {
+                using (var connection = new TestConnection(site.HttpClient.BaseAddress.Port))
+                {
+                    await connection.Send(
+                        "POST /ReadAndFlushEcho HTTP/1.1",
+                        $"Content-Length: {maxRequestSize}",
+                        "Host: localhost",
+                        "Connection: close",
+                        "",
+                        "");
+
+                    var disconnectAfter = random.Next(maxRequestSize);
+                    var data = new byte[blockSize];
+                    for (int i = 0; i < disconnectAfter / blockSize; i++)
+                    {
+                        await connection.Stream.WriteAsync(data);
+                    }
+                }
+            }
+
+            List<Task> tasks = new List<Task>();
+            for (int i = 0; i < 100; i++)
+            {
+                tasks.Add(Task.Run(RunRequests));
+            }
+
+            await Task.WhenAll(tasks);
+
+            StopServer();
+        }
+    }
+}

test/WebSites/InProcessWebSite/Startup.cs

@@ -324,6 +324,17 @@ namespace TestSite
                 result = await ctx.Request.Body.ReadAsync(readBuffer, 0, readBuffer.Length);
             }
         }
+        private async Task ReadAndFlushEcho(HttpContext ctx)
+        {
+            var readBuffer = new byte[4096];
+            var result = await ctx.Request.Body.ReadAsync(readBuffer, 0, readBuffer.Length);
+            while (result != 0)
+            {
+                await ctx.Response.WriteAsync(Encoding.UTF8.GetString(readBuffer, 0, result));
+                await ctx.Response.Body.FlushAsync();
+                result = await ctx.Request.Body.ReadAsync(readBuffer, 0, readBuffer.Length);
+            }
+        }
 
         private async Task ReadAndWriteEchoLines(HttpContext ctx)
         {

tools/SetupTestEnvironment.ps1

@@ -62,6 +62,8 @@ if (!(Test-Path $cdb))
 
 if ($Mode -eq "Setup")
 {
+    Move-Item $env:windir\System32\vsjitdebugger.exe $env:windir\System32\_vsjitdebugger.exe;
+
     Setup-appverif w3wp.exe
     Setup-appverif iisexpress.exe
 
@@ -70,6 +72,8 @@ if ($Mode -eq "Setup")
         New-Item -Path $werHive -Name LocalDumps
     }
 
+    New-ItemProperty $werHive -Name "DontShowUI" -Value 1 -PropertyType "DWORD" -Force;
+
     New-ItemProperty $ldHive -Name "DumpFolder" -Value $DumpFolder -PropertyType "ExpandString" -Force;
     New-ItemProperty $ldHive -Name "DumpCount" -Value 15 -PropertyType "DWORD" -Force;
     New-ItemProperty $ldHive -Name "DumpType" -Value 2 -PropertyType "DWORD" -Force;
@@ -79,8 +83,12 @@ if ($Mode -eq "Setup")
 
 if ($Mode -eq "Shutdown")
 {
+    Move-Item $env:windir\System32\_vsjitdebugger.exe $env:windir\System32\vsjitdebugger.exe;
+
     Remove-Item $ldHive -Recurse -Force
 
+    New-ItemProperty $werHive -Name "DontShowUI" -Value 0 -PropertyType "DWORD" -Force;
+
     Shutdown-appverif w3wp.exe
     Shutdown-appverif iisexpress.exe