Drain requests in native instead of managed

src/Servers/IIS/AspNetCoreModuleV2/CommonLib/application.h

@@ -7,6 +7,7 @@
 #include "iapplication.h"
 #include "ntassert.h"
 #include "SRWExclusiveLock.h"
+#include "SRWSharedLock.h"
 #include "exceptions.h"
 
 class APPLICATION : public IAPPLICATION
@@ -21,14 +22,17 @@ class APPLICATION : public IAPPLICATION
         _In_  IHttpContext       *pHttpContext,
         _Outptr_result_maybenull_ IREQUEST_HANDLER  **pRequestHandler) override
     {
-        TraceContextScope traceScope(pHttpContext->GetTraceContext());
         *pRequestHandler = nullptr;
 
+        SRWSharedLock stopLock(m_stateLock);
+
         if (m_fStopCalled)
         {
             return S_FALSE;
         }
 
+        TraceContextScope traceScope(pHttpContext->GetTraceContext());
+
         return CreateHandler(pHttpContext, pRequestHandler);
     }
 

src/Servers/IIS/AspNetCoreModuleV2/InProcessRequestHandler/inprocessapplication.cpp

@@ -26,7 +26,8 @@ IN_PROCESS_APPLICATION::IN_PROCESS_APPLICATION(
     m_Initialized(false),
     m_blockManagedCallbacks(true),
     m_waitForShutdown(true),
-    m_pConfig(std::move(pConfig))
+    m_pConfig(std::move(pConfig)),
+    m_requestCount(0)
 {
     DBG_ASSERT(m_pConfig);
 
@@ -67,6 +68,13 @@ IN_PROCESS_APPLICATION::StopClr()
         {
             shutdownHandler(m_ShutdownHandlerContext);
         }
+
+        auto requestCount = m_requestCount.load();
+        if (requestCount == 0)
+        {
+            LOG_INFO(L"Drained all requests, notifying managed.");
+            m_RequestsDrainedHandler(m_ShutdownHandlerContext);
+        }
     }
 
     // Signal shutdown
@@ -90,6 +98,7 @@ IN_PROCESS_APPLICATION::SetCallbackHandles(
     _In_ PFN_SHUTDOWN_HANDLER shutdown_handler,
     _In_ PFN_DISCONNECT_HANDLER disconnect_callback,
     _In_ PFN_ASYNC_COMPLETION_HANDLER async_completion_handler,
+    _In_ PFN_REQUESTS_DRAINED_HANDLER requestsDrainedHandler,
     _In_ VOID* pvRequstHandlerContext,
     _In_ VOID* pvShutdownHandlerContext
 )
@@ -102,6 +111,7 @@ IN_PROCESS_APPLICATION::SetCallbackHandles(
     m_ShutdownHandler = shutdown_handler;
     m_ShutdownHandlerContext = pvShutdownHandlerContext;
     m_AsyncCompletionHandler = async_completion_handler;
+    m_RequestsDrainedHandler = requestsDrainedHandler;
 
     m_blockManagedCallbacks = false;
     m_Initialized = true;
@@ -131,6 +141,12 @@ IN_PROCESS_APPLICATION::LoadManagedApplication()
         FALSE,    // not set
         nullptr)); // name
 
+    THROW_LAST_ERROR_IF_NULL(m_pRequestDrainEvent = CreateEvent(
+        nullptr,  // default security attributes
+        TRUE,     // manual reset event
+        FALSE,    // not set
+        nullptr)); // name
+
     LOG_INFO(L"Waiting for initialization");
 
     m_workerThread = std::thread([](std::unique_ptr<IN_PROCESS_APPLICATION, IAPPLICATION_DELETER> application)
@@ -167,7 +183,6 @@ IN_PROCESS_APPLICATION::LoadManagedApplication()
     return S_OK;
 }
 
-
 void
 IN_PROCESS_APPLICATION::ExecuteApplication()
 {
@@ -259,7 +274,7 @@ IN_PROCESS_APPLICATION::ExecuteApplication()
             if (m_waitForShutdown)
             {
                 const auto clrWaitResult = WaitForSingleObject(m_clrThread.native_handle(), m_pConfig->QueryShutdownTimeLimitInMS());
-                THROW_LAST_ERROR_IF(waitResult == WAIT_FAILED);
+                THROW_LAST_ERROR_IF(clrWaitResult == WAIT_FAILED);
 
                 clrThreadExited = clrWaitResult != WAIT_TIMEOUT;
             }
@@ -517,9 +532,23 @@ IN_PROCESS_APPLICATION::CreateHandler(
 {
     try
     {
+        DBG_ASSERT(!m_fStopCalled);
+        m_requestCount++;
         *pRequestHandler = new IN_PROCESS_HANDLER(::ReferenceApplication(this), pHttpContext, m_RequestHandler, m_RequestHandlerContext, m_DisconnectHandler, m_AsyncCompletionHandler);
     }
     CATCH_RETURN();
 
     return S_OK;
 }
+
+void
+IN_PROCESS_APPLICATION::HandleRequestCompletion()
+{
+    SRWSharedLock lock(m_stateLock);
+    auto requestCount = m_requestCount--;
+    if (m_fStopCalled && requestCount == 0)
+    {
+        LOG_INFO(L"Drained all requests, notifying managed.");
+        m_RequestsDrainedHandler(m_ShutdownHandlerContext);
+    }
+}

src/Servers/IIS/AspNetCoreModuleV2/InProcessRequestHandler/inprocessapplication.h

@@ -13,6 +13,7 @@ typedef REQUEST_NOTIFICATION_STATUS(WINAPI * PFN_REQUEST_HANDLER) (IN_PROCESS_HA
 typedef VOID(WINAPI * PFN_DISCONNECT_HANDLER) (void *pvManagedHttpContext);
 typedef BOOL(WINAPI * PFN_SHUTDOWN_HANDLER) (void* pvShutdownHandlerContext);
 typedef REQUEST_NOTIFICATION_STATUS(WINAPI * PFN_ASYNC_COMPLETION_HANDLER)(void *pvManagedHttpContext, HRESULT hrCompletionStatus, DWORD cbCompletion);
+typedef void(WINAPI * PFN_REQUESTS_DRAINED_HANDLER) (void* pvShutdownHandlerContext);
 
 class IN_PROCESS_APPLICATION : public InProcessApplicationBase
 {
@@ -36,6 +37,7 @@ class IN_PROCESS_APPLICATION : public InProcessApplicationBase
         _In_ PFN_SHUTDOWN_HANDLER shutdown_callback,
         _In_ PFN_DISCONNECT_HANDLER disconnect_callback,
         _In_ PFN_ASYNC_COMPLETION_HANDLER managed_context_callback,
+        _In_ PFN_REQUESTS_DRAINED_HANDLER requestsDrainedHandler,
         _In_ VOID* pvRequstHandlerContext,
         _In_ VOID* pvShutdownHandlerContext
     );
@@ -47,6 +49,8 @@ class IN_PROCESS_APPLICATION : public InProcessApplicationBase
         _Out_ IREQUEST_HANDLER   **pRequestHandler)
     override;
 
+    void HandleRequestCompletion();
+
     // Executes the .NET Core process
     void
     ExecuteApplication();
@@ -62,6 +66,8 @@ class IN_PROCESS_APPLICATION : public InProcessApplicationBase
     StopIncomingRequests()
     {
         QueueStop();
+
+        LOG_INFOF(L"Waiting for %d requests to drain", m_requestCount.load());
     }
 
     void
@@ -144,6 +150,8 @@ class IN_PROCESS_APPLICATION : public InProcessApplicationBase
     // The event that gets triggered when worker thread should exit
     HandleWrapper<NullHandleTraits> m_pShutdownEvent;
 
+    HandleWrapper<NullHandleTraits> m_pRequestDrainEvent;
+
     // The request handler callback from managed code
     PFN_REQUEST_HANDLER             m_RequestHandler;
     VOID*                           m_RequestHandlerContext;
@@ -154,12 +162,14 @@ class IN_PROCESS_APPLICATION : public InProcessApplicationBase
 
     PFN_ASYNC_COMPLETION_HANDLER    m_AsyncCompletionHandler;
     PFN_DISCONNECT_HANDLER          m_DisconnectHandler;
+    PFN_REQUESTS_DRAINED_HANDLER    m_RequestsDrainedHandler;
 
     std::wstring                    m_dotnetExeKnownLocation;
 
     std::atomic_bool                m_blockManagedCallbacks;
     bool                            m_Initialized;
     bool                            m_waitForShutdown;
+    std::atomic<int>                m_requestCount;
 
     std::unique_ptr<InProcessOptions> m_pConfig;
 

src/Servers/IIS/AspNetCoreModuleV2/InProcessRequestHandler/inprocesshandler.cpp

@@ -58,7 +58,6 @@ IN_PROCESS_HANDLER::AsyncCompletion(
     HRESULT     hrCompletionStatus
 )
 {
-
     ::RaiseEvent<ANCMEvents::ANCM_INPROC_ASYNC_COMPLETION_START>(m_pW3Context, nullptr);
 
     if (m_fManagedRequestComplete)

src/Servers/IIS/AspNetCoreModuleV2/InProcessRequestHandler/inprocesshandler.h

@@ -21,7 +21,10 @@ class IN_PROCESS_HANDLER : public REQUEST_HANDLER
         _In_ PFN_DISCONNECT_HANDLER m_DisconnectHandler,
         _In_ PFN_ASYNC_COMPLETION_HANDLER pAsyncCompletion);
 
-    ~IN_PROCESS_HANDLER() override = default;
+    ~IN_PROCESS_HANDLER()
+    {
+        m_pApplication->HandleRequestCompletion();
+    }
 
     __override
     REQUEST_NOTIFICATION_STATUS

src/Servers/IIS/AspNetCoreModuleV2/InProcessRequestHandler/managedexports.cpp

@@ -18,6 +18,7 @@ register_callbacks(
     _In_ PFN_SHUTDOWN_HANDLER shutdown_handler,
     _In_ PFN_DISCONNECT_HANDLER disconnect_handler,
     _In_ PFN_ASYNC_COMPLETION_HANDLER async_completion_handler,
+    _In_ PFN_REQUESTS_DRAINED_HANDLER requestsDrainedHandler,
     _In_ VOID* pvRequstHandlerContext,
     _In_ VOID* pvShutdownHandlerContext
 )
@@ -32,6 +33,7 @@ register_callbacks(
         shutdown_handler,
         disconnect_handler,
         async_completion_handler,
+        requestsDrainedHandler,
         pvRequstHandlerContext,
         pvShutdownHandlerContext
     );

src/Servers/IIS/IIS/src/Core/IISHttpContext.cs

@@ -556,7 +556,6 @@ private async Task HandleRequest()
                 // Post completion after completing the request to resume the state machine
                 PostCompletion(ConvertRequestCompletionResults(successfulRequest));
 
-                Server.DecrementRequests();
 
                 // Dispose the context
                 Dispose();

src/Servers/IIS/IIS/src/Core/IISHttpServer.cs

@@ -1,3 +1,5 @@
+// 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.Buffers;
@@ -23,6 +25,7 @@ internal class IISHttpServer : IServer
         private static readonly NativeMethods.PFN_SHUTDOWN_HANDLER _shutdownHandler = HandleShutdown;
         private static readonly NativeMethods.PFN_DISCONNECT_HANDLER _onDisconnect = OnDisconnect;
         private static readonly NativeMethods.PFN_ASYNC_COMPLETION _onAsyncCompletion = OnAsyncCompletion;
+        private static readonly NativeMethods.PFN_REQUESTS_DRAINED_HANDLER _requestsDrainedHandler = OnRequestsDrained;
 
         private IISContextFactory _iisContextFactory;
         private readonly MemoryPool<byte> _memoryPool = new SlabMemoryPool();
@@ -33,11 +36,9 @@ internal class IISHttpServer : IServer
         private readonly IISNativeApplication _nativeApplication;
         private readonly ServerAddressesFeature _serverAddressesFeature;
 
-        private volatile int _stopping;
-        private bool Stopping => _stopping == 1;
-        private int _outstandingRequests;
         private readonly TaskCompletionSource<object> _shutdownSignal = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
         private bool? _websocketAvailable;
+        private CancellationTokenRegistration _cancellationTokenRegistration;
 
         public IFeatureCollection Features { get; } = new FeatureCollection();
 
@@ -86,7 +87,7 @@ public Task StartAsync<TContext>(IHttpApplication<TContext> application, Cancell
             _httpServerHandle = GCHandle.Alloc(this);
 
             _iisContextFactory = new IISContextFactory<TContext>(_memoryPool, application, _options, this, _logger);
-            _nativeApplication.RegisterCallbacks(_requestHandler, _shutdownHandler, _onDisconnect, _onAsyncCompletion, (IntPtr)_httpServerHandle, (IntPtr)_httpServerHandle);
+            _nativeApplication.RegisterCallbacks(_requestHandler, _shutdownHandler, _onDisconnect, _onAsyncCompletion, _requestsDrainedHandler, (IntPtr)_httpServerHandle, (IntPtr)_httpServerHandle);
 
             _serverAddressesFeature.Addresses = _options.ServerAddresses;
 
@@ -95,50 +96,18 @@ public Task StartAsync<TContext>(IHttpApplication<TContext> application, Cancell
 
         public Task StopAsync(CancellationToken cancellationToken)
         {
-            void RegisterCancelation()
-            {
-                cancellationToken.Register(() =>
-                {
-                    _nativeApplication.StopCallsIntoManaged();
-                    _shutdownSignal.TrySetResult(null);
-                });
-            }
-            if (Interlocked.Exchange(ref _stopping, 1) == 1)
-            {
-                RegisterCancelation();
-
-                return _shutdownSignal.Task;
-            }
-
-            // First call back into native saying "DON'T SEND ME ANY MORE REQUESTS"
             _nativeApplication.StopIncomingRequests();
-
-            try
+            _cancellationTokenRegistration = cancellationToken.Register((shutdownSignal) =>
             {
-                // Wait for active requests to drain
-                if (_outstandingRequests > 0)
-                {
-                    RegisterCancelation();
-                }
-                else
-                {
-                    // We have drained all requests. Block any callbacks into managed at this point.
-                    _nativeApplication.StopCallsIntoManaged();
-                    _shutdownSignal.TrySetResult(null);
-                }
-            }
-            catch (Exception ex)
-            {
-                _shutdownSignal.TrySetException(ex);
-            }
+                ((TaskCompletionSource<object>)shutdownSignal).TrySetResult(null);
+            },
+            _shutdownSignal);
 
             return _shutdownSignal.Task;
         }
 
         public void Dispose()
         {
-            _stopping = 1;
-
             // Block any more calls into managed from native as we are unloading.
             _nativeApplication.StopCallsIntoManaged();
             _shutdownSignal.TrySetResult(null);
@@ -152,29 +121,13 @@ public void Dispose()
             _nativeApplication.Dispose();
         }
 
-        private void IncrementRequests()
-        {
-            Interlocked.Increment(ref _outstandingRequests);
-        }
-
-        internal void DecrementRequests()
-        {
-            if (Interlocked.Decrement(ref _outstandingRequests) == 0 && Stopping)
-            {
-                // All requests have been drained.
-                _nativeApplication.StopCallsIntoManaged();
-                _shutdownSignal.TrySetResult(null);
-            }
-        }
-
         private static NativeMethods.REQUEST_NOTIFICATION_STATUS HandleRequest(IntPtr pInProcessHandler, IntPtr pvRequestContext)
         {
             IISHttpServer server = null;
             try
             {
                 // Unwrap the server so we can create an http context and process the request
                 server = (IISHttpServer)GCHandle.FromIntPtr(pvRequestContext).Target;
-                server.IncrementRequests();
 
                 var context = server._iisContextFactory.CreateHttpContext(pInProcessHandler);
 
@@ -205,7 +158,6 @@ private static bool HandleShutdown(IntPtr pvRequestContext)
             return true;
         }
 
-
         private static void OnDisconnect(IntPtr pvManagedHttpContext)
         {
             IISHttpContext context = null;
@@ -237,6 +189,23 @@ private static NativeMethods.REQUEST_NOTIFICATION_STATUS OnAsyncCompletion(IntPt
             }
         }
 
+        private static void OnRequestsDrained(IntPtr serverContext)
+        {
+            IISHttpServer server = null;
+            try
+            {
+                server = (IISHttpServer)GCHandle.FromIntPtr(serverContext).Target;
+
+                server._nativeApplication.StopCallsIntoManaged();
+                server._shutdownSignal.TrySetResult(null);
+                server._cancellationTokenRegistration.Dispose();
+            }
+            catch (Exception ex)
+            {
+                server?._logger.LogError(0, ex, $"Unexpected exception in {nameof(IISHttpServer)}.{nameof(OnRequestsDrained)}.");
+            }
+        }
+
         private class IISContextFactory<T> : IISContextFactory
         {
             private readonly IHttpApplication<T> _application;

src/Servers/IIS/IIS/src/Core/IISNativeApplication.cs

@@ -30,6 +30,7 @@ public void RegisterCallbacks(
             NativeMethods.PFN_SHUTDOWN_HANDLER shutdownHandler,
             NativeMethods.PFN_DISCONNECT_HANDLER disconnectHandler,
             NativeMethods.PFN_ASYNC_COMPLETION onAsyncCompletion,
+            NativeMethods.PFN_REQUESTS_DRAINED_HANDLER requestDrainedHandler,
             IntPtr requestContext,
             IntPtr shutdownContext)
         {
@@ -39,6 +40,7 @@ public void RegisterCallbacks(
                 shutdownHandler,
                 disconnectHandler,
                 onAsyncCompletion,
+                requestDrainedHandler,
                 requestContext,
                 shutdownContext);
         }