Make CTRL signal able to take multiple contexts (#1172)

Update the CTRL signal handler to be able to fan out to multiple contexts, and make all cases of using sub-contexts both connected to the signal handler and better handle cancellation.

Additionally, improve the COM API install method to properly cancel any in progress task rather than just terminating the context.

.github/actions/spelling/allow.txt

@@ -424,6 +424,7 @@ STRINGIFY
 STRINGIZE
 stringstream
 strstr
+subcontext
 SUBLANG
 subresource
 subselect

src/AppInstallerCLICore/ExecutionContext.cpp

@@ -10,54 +10,101 @@ namespace AppInstaller::CLI::Execution
 
     namespace
     {
-        // The context that will receive CTRL signals
-        Context* s_contextForCtrlHandler = nullptr;
-
-        BOOL WINAPI CtrlHandlerForContext(DWORD ctrlType)
+        // Type to contain the CTRL signal handler.
+        struct CtrlHandler
         {
-            AICLI_LOG(CLI, Info, << "Got CTRL type: " << ctrlType);
+            static CtrlHandler& Instance()
+            {
+                static CtrlHandler s_instance;
+                return s_instance;
+            }
 
-            // Won't save us from every crash, but a few more than direct access.
-            Context* context = s_contextForCtrlHandler;
-            if (!context)
+            void AddContext(Context* context)
             {
-                return FALSE;
+                std::lock_guard<std::mutex> lock{ m_contextsLock };
+
+                auto itr = std::find(m_contexts.begin(), m_contexts.end(), context);
+                THROW_HR_IF(E_NOT_VALID_STATE, itr != m_contexts.end());
+                m_contexts.push_back(context);
             }
 
-            switch (ctrlType)
+            void RemoveContext(Context* context)
             {
-            case CTRL_C_EVENT:
-            case CTRL_BREAK_EVENT:
-                context->Terminate(APPINSTALLER_CLI_ERROR_CTRL_SIGNAL_RECEIVED);
-                context->Reporter.CancelInProgressTask(false);
-                return TRUE;
-                // According to MSDN, we should never receive these due to having gdi32/user32 loaded in our process.
-                // But handle them as a force terminate anyway.
-            case CTRL_CLOSE_EVENT:
-            case CTRL_LOGOFF_EVENT:
-            case CTRL_SHUTDOWN_EVENT:
-                context->Terminate(APPINSTALLER_CLI_ERROR_CTRL_SIGNAL_RECEIVED);
-                context->Reporter.CancelInProgressTask(true);
+                std::lock_guard<std::mutex> lock{ m_contextsLock };
+
+                auto itr = std::find(m_contexts.begin(), m_contexts.end(), context);
+                THROW_HR_IF(E_NOT_VALID_STATE, itr == m_contexts.end());
+                m_contexts.erase(itr);
+            }
+
+        private:
+            CtrlHandler()
+            {
+                LOG_IF_WIN32_BOOL_FALSE(SetConsoleCtrlHandler(StaticCtrlHandlerFunction, TRUE));
+            }
+
+            static BOOL WINAPI StaticCtrlHandlerFunction(DWORD ctrlType)
+            {
+                return Instance().CtrlHandlerFunction(ctrlType);
+            }
+
+            BOOL CtrlHandlerFunction(DWORD ctrlType)
+            {
+                switch (ctrlType)
+                {
+                case CTRL_C_EVENT:
+                case CTRL_BREAK_EVENT:
+                    return TerminateContexts(ctrlType, false);
+                    // According to MSDN, we should never receive these due to having gdi32/user32 loaded in our process.
+                    // But handle them as a force terminate anyway.
+                case CTRL_CLOSE_EVENT:
+                case CTRL_LOGOFF_EVENT:
+                case CTRL_SHUTDOWN_EVENT:
+                    return TerminateContexts(ctrlType, true);
+                default:
+                    return FALSE;
+                }
+            }
+
+            // Terminates the currently attached contexts.
+            // Returns FALSE if no contexts attached; TRUE otherwise.
+            BOOL TerminateContexts(DWORD ctrlType, bool force)
+            {
+                if (m_contexts.empty())
+                {
+                    return FALSE;
+                }
+
+                {
+                    std::lock_guard<std::mutex> lock{ m_contextsLock };
+
+                    // TODO: Move this to be logged per active context when we have thread static globals
+                    AICLI_LOG(CLI, Info, << "Got CTRL type: " << ctrlType);
+
+                    for (auto& context : m_contexts)
+                    {
+                        context->Cancel(true, force);
+                    }
+                }
+
                 return TRUE;
-            default:
-                return FALSE;
             }
-        }
 
-        void SetCtrlHandlerContext(Context* context)
+            std::mutex m_contextsLock;
+            std::vector<Context*> m_contexts;
+        };
+
+        void SetCtrlHandlerContext(bool add, Context* context)
         {
-            // Only one is allowed right now.
-            THROW_HR_IF(E_UNEXPECTED, s_contextForCtrlHandler != nullptr && context != nullptr);
+            THROW_HR_IF(E_POINTER, context == nullptr);
 
-            if (context == nullptr)
+            if (add)
             {
-                LOG_IF_WIN32_BOOL_FALSE(SetConsoleCtrlHandler(CtrlHandlerForContext, FALSE));
-                s_contextForCtrlHandler = nullptr;
+                CtrlHandler::Instance().AddContext(context);
             }
             else
             {
-                s_contextForCtrlHandler = context;
-                LOG_IF_WIN32_BOOL_FALSE(SetConsoleCtrlHandler(CtrlHandlerForContext, TRUE));
+                CtrlHandler::Instance().RemoveContext(context);
             }
         }
     }
@@ -74,12 +121,17 @@ namespace AppInstaller::CLI::Execution
     {
         auto clone = std::make_unique<Context>(Reporter);
         clone->m_flags = m_flags;
+        // If the parent is hooked up to the CTRL signal, have the clone be as well
+        if (m_disableCtrlHandlerOnExit)
+        {
+            clone->EnableCtrlHandler();
+        }
         return clone;
     }
 
     void Context::EnableCtrlHandler(bool enabled)
     {
-        SetCtrlHandlerContext(enabled ? this : nullptr);
+        SetCtrlHandlerContext(enabled, this);
         m_disableCtrlHandlerOnExit = enabled;
     }
 
@@ -126,9 +178,15 @@ namespace AppInstaller::CLI::Execution
         m_isTerminated = true;
         m_terminationHR = hr;
     }
-
+
+    void Context::Cancel(bool exitIfStuck, bool bypassUser)
+    {
+        Terminate(exitIfStuck ? APPINSTALLER_CLI_ERROR_CTRL_SIGNAL_RECEIVED : E_ABORT);
+        Reporter.CancelInProgressTask(bypassUser);
+    }
+
     void Context::SetExecutionStage(Workflow::ExecutionStage stage, bool allowBackward)
-    {     
+    {
         if (m_executionStage == stage)
         {
             return;

src/AppInstallerCLICore/ExecutionContext.h

@@ -82,6 +82,11 @@ namespace AppInstaller::CLI::Execution
         // Set the context to the terminated state.
         void Terminate(HRESULT hr, std::string_view file = {}, size_t line = {});
 
+        // Cancel the context; this terminates it as well as informing any in progress task to stop cooperatively.
+        // Multiple attempts with exitIfStuck == true may cause the process to simply exit.
+        // The bypassUser indicates whether the user should be asked for cancellation (does not currently have any effect).
+        void Cancel(bool exitIfStuck = false, bool bypassUser = false);
+
         // Gets context flags
         ContextFlag GetFlags() const
         {

src/AppInstallerCLICore/Workflows/ImportExportFlow.cpp

@@ -284,7 +284,13 @@ namespace AppInstaller::CLI::Workflow
 
                 if (searchContext.IsTerminated())
                 {
-                    if (searchContext.GetTerminationHR() == APPINSTALLER_CLI_ERROR_UPDATE_NOT_APPLICABLE)
+                    if (context.IsTerminated() && context.GetTerminationHR() == E_ABORT)
+                    {
+                        // This means that the subcontext being terminated is due to an overall abort
+                        context.Reporter.Info() << Resource::String::Cancelled << std::endl;
+                        return;
+                    }
+                    else if (searchContext.GetTerminationHR() == APPINSTALLER_CLI_ERROR_UPDATE_NOT_APPLICABLE)
                     {
                         AICLI_LOG(CLI, Info, << "Package is already installed: [" << packageRequest.Id << "]");
                         context.Reporter.Info() << Resource::String::ImportPackageAlreadyInstalled << ' ' << packageRequest.Id << std::endl;

src/AppInstallerCLICore/Workflows/InstallFlow.cpp

@@ -432,6 +432,13 @@ namespace AppInstaller::CLI::Workflow
             installContext << InstallPackageVersion;
             if (installContext.IsTerminated())
             {
+                if (context.IsTerminated() && context.GetTerminationHR() == E_ABORT)
+                {
+                    // This means that the subcontext being terminated is due to an overall abort
+                    context.Reporter.Info() << Resource::String::Cancelled << std::endl;
+                    return;
+                }
+
                 allSucceeded = false;
             }
         }

src/AppInstallerCLICore/Workflows/UpdateFlow.cpp

@@ -121,6 +121,12 @@ namespace AppInstaller::CLI::Workflow
             {
                 updateAllHasFailure = true;
             }
+
+            if (context.IsTerminated() && context.GetTerminationHR() == E_ABORT)
+            {
+                context.Reporter.Info() << Resource::String::Cancelled << std::endl;
+                return;
+            }
         }
 
         if (!updateAllFoundUpdate)

src/Microsoft.Management.Deployment/PackageManager.cpp

@@ -299,7 +299,7 @@ namespace winrt::Microsoft::Management::Deployment::implementation
 
         cancellationToken.callback([&context]
             {
-                context.Terminate(APPINSTALLER_CLI_ERROR_CTRL_SIGNAL_RECEIVED);
+                context.Cancel(false, true);
             });
         // Wait for the execute operation to finish. 
         // The cancellation of the AsyncOperation triggers Terminate which causes the executeOperation to end.