winget-cli

commit 30b652b6bb96116d61ebf05f9589d0c8bdc1a3ec
parent 712486a0eb007c3f277c4e982f41089dd9372049
Author: JohnMcPMS <johnmcp@microsoft.com>
Date:   Fri, 27 Mar 2020 14:10:25 -0700

Add CTRL handler to enable cooperative cancel (#74)


Diffstat:
M
src/AppInstallerCLICore/Core.cpp
 | 
1
+
M
src/AppInstallerCLICore/ExecutionReporter.cpp
 | 
73
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
M
src/AppInstallerCLICore/ExecutionReporter.h
 | 
19
+++++++++++++++++--
M
src/AppInstallerCommonCore/Public/AppInstallerProgress.h
 | 
1
-

4 files changed, 91 insertions(+), 3 deletions(-)
diff --git a/src/AppInstallerCLICore/Core.cpp b/src/AppInstallerCLICore/Core.cpp
@@ -54,6 +54,7 @@ namespace AppInstaller::CLI
         Logging::Telemetry().LogStartup();
 
         Execution::Context context{ std::cout, std::cin };
+        context.Reporter.EnableCtrlHandler();
 
         // Convert incoming wide char args to UTF8
         std::vector<std::string> utf8Args;
diff --git a/src/AppInstallerCLICore/ExecutionReporter.cpp b/src/AppInstallerCLICore/ExecutionReporter.cpp
@@ -6,6 +6,50 @@
 
 namespace AppInstaller::CLI::Execution
 {
+    namespace
+    {
+        // The reporter that will receive CTRL signals
+        Reporter* s_reporterForCtrlHandler = nullptr;
+
+        BOOL WINAPI CtrlHandlerForReporter(DWORD ctrlType)
+        {
+            switch (ctrlType)
+            {
+            case CTRL_C_EVENT:
+            case CTRL_BREAK_EVENT:
+                s_reporterForCtrlHandler->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:
+                s_reporterForCtrlHandler->CancelInProgressTask(true);
+                return TRUE;
+            default:
+                AICLI_LOG(CLI, Info, << "Got unrecognized CTRL type: " << ctrlType);
+                return FALSE;
+            }
+        }
+
+        void SetCtrlHandlerReporter(Reporter* reporter)
+        {
+            // Only one is allowed right now.
+            THROW_HR_IF(E_UNEXPECTED, s_reporterForCtrlHandler != nullptr && reporter != nullptr);
+
+            if (reporter == nullptr)
+            {
+                LOG_IF_WIN32_BOOL_FALSE(SetConsoleCtrlHandler(CtrlHandlerForReporter, FALSE));
+                s_reporterForCtrlHandler = nullptr;
+            }
+            else
+            {
+                s_reporterForCtrlHandler = reporter;
+                LOG_IF_WIN32_BOOL_FALSE(SetConsoleCtrlHandler(CtrlHandlerForReporter, TRUE));
+            }
+        }
+    }
+
     namespace details
     {
         void IndefiniteSpinner::ShowSpinner()
@@ -119,6 +163,11 @@ namespace AppInstaller::CLI::Execution
         {
             m_out << VirtualTerminal::TextFormat::Default;
         }
+
+        if (m_disableCtrlHandlerOnExit)
+        {
+            EnableCtrlHandler(false);
+        }
     }
 
     Reporter::OutputStream Reporter::GetOutputStream(Level level)
@@ -186,4 +235,28 @@ namespace AppInstaller::CLI::Execution
         ShowIndefiniteProgress(false);
         ShowProgress(true, (maximum ? static_cast<uint64_t>((static_cast<double>(current) / maximum) * 100) : current));
     }
+
+    void Reporter::EnableCtrlHandler(bool enabled)
+    {
+        SetCtrlHandlerReporter(enabled ? this : nullptr);
+        m_disableCtrlHandlerOnExit = enabled;
+    }
+
+    void Reporter::SetProgressCallback(ProgressCallback* callback)
+    {
+        auto lock = m_progressCallbackLock.lock_exclusive();
+        m_progressCallback = callback;
+    }
+
+    void Reporter::CancelInProgressTask(bool force)
+    {
+        // TODO: Maybe ask the user if they really want to cancel?
+        UNREFERENCED_PARAMETER(force);
+        auto lock = m_progressCallbackLock.lock_shared();
+        ProgressCallback* callback = m_progressCallback.load();
+        if (callback)
+        {
+            callback->Cancel();
+        }
+    }
 }
diff --git a/src/AppInstallerCLICore/ExecutionReporter.h b/src/AppInstallerCLICore/ExecutionReporter.h
@@ -50,8 +50,8 @@ namespace AppInstaller::CLI::Execution
         };
     }
 
-    // WorkflowReporter should be the central place to show workflow status to user.
-    // Todo: need to implement actual console output to show color, progress bar, etc
+    // Reporter should be the central place to show workflow status to user.
+    // Todo: need to implement actual console output to show progress bar, etc
     struct Reporter : public IProgressSink
     {
         enum class Level
@@ -134,21 +134,36 @@ namespace AppInstaller::CLI::Execution
         auto ExecuteWithProgress(F&& f)
         {
             ProgressCallback callback(this);
+            SetProgressCallback(&callback);
             ShowIndefiniteProgress(true);
 
             auto hideProgress = wil::scope_exit([this]()
                 {
+                    SetProgressCallback(nullptr);
                     ShowIndefiniteProgress(false);
                     ShowProgress(false, 0);
                 });
             return f(callback);
         }
 
+        // Enables reception of CTRL signals.
+        // Only one reporter can be enabled to handle CTRL signals at a time.
+        void EnableCtrlHandler(bool enabled = true);
+
+        // Sets the in progress callback.
+        void SetProgressCallback(ProgressCallback* callback);
+
+        // Cancels the in progress task.
+        void CancelInProgressTask(bool force);
+
     private:
         std::ostream& m_out;
         std::istream& m_in;
         VirtualTerminal::ConsoleModeRestore m_consoleMode;
         details::IndefiniteSpinner m_spinner;
         details::ProgressBar m_progressBar;
+        DestructionToken m_disableCtrlHandlerOnExit = false;
+        wil::srwlock m_progressCallbackLock;
+        std::atomic<ProgressCallback*> m_progressCallback;
     };
 }
diff --git a/src/AppInstallerCommonCore/Public/AppInstallerProgress.h b/src/AppInstallerCommonCore/Public/AppInstallerProgress.h
@@ -31,7 +31,6 @@ namespace AppInstaller
     // of cancel state.
     struct IProgressSink
     {
-
         // Called as progress is made.
         // If maximum is 0, the maximum is unknown.
         virtual void OnProgress(uint64_t current, uint64_t maximum, ProgressType type) = 0;