winget-cli

ShutdownMonitoring.cpp (11580B)

---

 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 #include "pch.h"
 #include "Public/ShutdownMonitoring.h"
 #include <AppInstallerErrors.h>
 #include <AppInstallerLogging.h>
 #include <AppInstallerRuntime.h>
 #include <winget/COMStaticStorage.h>
 
 namespace AppInstaller::ShutdownMonitoring
 {
     std::shared_ptr<TerminationSignalHandler> TerminationSignalHandler::Instance()
     {
         struct Singleton : public WinRT::COMStaticStorageBase<TerminationSignalHandler>
         {
             Singleton() : COMStaticStorageBase(L"WindowsPackageManager.TerminationSignalHandler") {}
         };
 
         static Singleton s_instance;
         return s_instance.Get();
     }
 
     void TerminationSignalHandler::AddListener(ICancellable* cancellable)
     {
         std::lock_guard<std::mutex> lock{ m_listenersLock };
 
         auto itr = std::find(m_listeners.begin(), m_listeners.end(), cancellable);
         THROW_HR_IF(E_NOT_VALID_STATE, itr != m_listeners.end());
         m_listeners.push_back(cancellable);
     }
 
     void TerminationSignalHandler::RemoveListener(ICancellable* cancellable)
     {
         std::lock_guard<std::mutex> lock{ m_listenersLock };
 
         auto itr = std::find(m_listeners.begin(), m_listeners.end(), cancellable);
         if (itr == m_listeners.end())
         {
             AICLI_LOG(CLI, Warning, << "TerminationSignalHandler::RemoveListener did not find requested object");
         }
         else
         {
             m_listeners.erase(itr);
         }
     }
 
     void TerminationSignalHandler::EnableListener(bool enabled, ICancellable* cancellable)
     {
         if (enabled)
         {
             Instance()->AddListener(cancellable);
         }
         else
         {
             Instance()->RemoveListener(cancellable);
         }
     }
 
 #ifndef AICLI_DISABLE_TEST_HOOKS
     HWND TerminationSignalHandler::GetWindowHandle() const
     {
         return m_windowHandle.get();
     }
 
     bool TerminationSignalHandler::WaitForAppShutdownEvent() const
     {
         return m_appShutdownEvent.wait(60000);
     }
 #endif
 
     TerminationSignalHandler::TerminationSignalHandler()
     {
 #ifndef AICLI_DISABLE_TEST_HOOKS
         m_appShutdownEvent.create();
 #endif
 
         if (Runtime::IsRunningInPackagedContext())
         {
             // Create package update listener
             m_catalog = winrt::Windows::ApplicationModel::PackageCatalog::OpenForCurrentPackage();
             m_updatingEvent = m_catalog.PackageUpdating(
                 winrt::auto_revoke, [this](winrt::Windows::ApplicationModel::PackageCatalog, winrt::Windows::ApplicationModel::PackageUpdatingEventArgs)
                 {
                     this->StartAppShutdown();
                 });
         }
 
         // Create message only window.
         m_messageQueueReady.create();
         m_windowThread = std::thread(&TerminationSignalHandler::CreateWindowAndStartMessageLoop, this);
         if (!m_messageQueueReady.wait(100))
         {
             AICLI_LOG(CLI, Warning, << "Timeout creating winget window");
         }
 
         // Set up ctrl-c handler.
         LOG_IF_WIN32_BOOL_FALSE(SetConsoleCtrlHandler(StaticCtrlHandlerFunction, TRUE));
     }
 
     TerminationSignalHandler::~TerminationSignalHandler()
     {
         // std::thread requires that any managed thread (joinable) be joined or detached before destructing
         if (m_windowThread.joinable())
         {
             m_windowThread.detach();
         }
     }
 
     void TerminationSignalHandler::StartAppShutdown()
     {
         AICLI_LOG(CLI, Info, << "Initiating shutdown procedure");
 
 #ifndef AICLI_DISABLE_TEST_HOOKS
         m_appShutdownEvent.SetEvent();
 #endif
 
         // Lifetime manager sends CTRL-C after the WM_QUERYENDSESSION is processed.
         // If we disable the CTRL-C handler, the default handler will kill us.
         InformListeners(CancelReason::AppShutdown, true);
     }
 
     BOOL WINAPI TerminationSignalHandler::StaticCtrlHandlerFunction(DWORD ctrlType)
     {
         return Instance()->CtrlHandlerFunction(ctrlType);
     }
 
     LRESULT WINAPI TerminationSignalHandler::WindowMessageProcedure(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
     {
         switch (uMsg)
         {
         case WM_QUERYENDSESSION:
             AICLI_LOG(CLI, Verbose, << "Received WM_QUERYENDSESSION");
             Instance()->StartAppShutdown();
             return TRUE;
         case WM_ENDSESSION:
         case WM_CLOSE:
             AICLI_LOG(CLI, Verbose, << "Received window message type: " << uMsg);
             // We delay as long as needed during the WM_ENDSESSION as we will be terminated on return.
             ServerShutdownSynchronization::WaitForShutdown();
             DestroyWindow(hWnd);
             break;
         case WM_DESTROY:
             PostQuitMessage(0);
             break;
         default:
             return DefWindowProc(hWnd, uMsg, wParam, lParam);
         }
         return FALSE;
     }
 
     BOOL TerminationSignalHandler::CtrlHandlerFunction(DWORD ctrlType)
     {
         // TODO: Move this to be logged per active context when we have thread static globals
         AICLI_LOG(CLI, Info, << "Got CTRL type: " << ctrlType);
 
         switch (ctrlType)
         {
         case CTRL_C_EVENT:
         case CTRL_BREAK_EVENT:
             return InformListeners(CancelReason::CtrlCSignal, 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 InformListeners(CancelReason::CtrlCSignal, true);
         default:
             return FALSE;
         }
     }
 
     // Terminates the currently attached contexts.
     // Returns FALSE if no contexts attached; TRUE otherwise.
     BOOL TerminationSignalHandler::InformListeners(CancelReason reason, bool force)
     {
         std::lock_guard<std::mutex> lock{ m_listenersLock };
 
         if (m_listeners.empty())
         {
             return FALSE;
         }
 
         for (auto& listener : m_listeners)
         {
             listener->Cancel(reason, force);
         }
 
         return TRUE;
     }
 
     void TerminationSignalHandler::CreateWindowAndStartMessageLoop()
     {
         PCWSTR windowClass = L"wingetWindow";
         HINSTANCE hInstance = GetModuleHandle(NULL);
         if (hInstance == NULL)
         {
             LOG_LAST_ERROR_MSG("Failed getting module handle");
             return;
         }
 
         WNDCLASSEX wcex = {};
         wcex.cbSize = sizeof(wcex);
 
         wcex.style = CS_NOCLOSE;
         wcex.lpfnWndProc = TerminationSignalHandler::WindowMessageProcedure;
         wcex.cbClsExtra = 0;
         wcex.cbWndExtra = 0;
         wcex.hInstance = hInstance;
         wcex.lpszClassName = windowClass;
 
         if (!RegisterClassEx(&wcex))
         {
             LOG_LAST_ERROR_MSG("Failed registering window class");
             return;
         }
 
         m_windowHandle = wil::unique_hwnd(CreateWindow(
             windowClass,
             L"WingetMessageOnlyWindow",
             WS_OVERLAPPEDWINDOW,
             0, /* x */
             0, /* y */
             0, /* nWidth */
             0, /* nHeight */
             NULL, /* hWndParent */
             NULL, /* hMenu */
             hInstance,
             NULL)); /* lpParam */
 
         if (m_windowHandle == nullptr)
         {
             LOG_LAST_ERROR_MSG("Failed creating window");
             return;
         }
 
         ShowWindow(m_windowHandle.get(), SW_HIDE);
 
         // Force message queue to be created.
         MSG msg;
         PeekMessage(&msg, NULL, WM_USER, WM_USER, PM_NOREMOVE);
         m_messageQueueReady.SetEvent();
 
         // Message loop
         BOOL getMessageResult;
         while ((getMessageResult = GetMessage(&msg, m_windowHandle.get(), 0, 0)) != 0)
         {
             if (getMessageResult == -1)
             {
                 LOG_LAST_ERROR();
             }
             else
             {
                 DispatchMessage(&msg);
             }
         }
     }
 
     void ServerShutdownSynchronization::Initialize(ShutdownCompleteCallback callback)
     {
         Instance().m_callback = callback;
     }
 
     void ServerShutdownSynchronization::AddComponent(const ComponentSystem& component)
     {
         ServerShutdownSynchronization& instance = Instance();
         std::lock_guard<std::mutex> lock{ instance.m_componentsLock };
 
         for (const auto& item : instance.m_components)
         {
             if (item.BlockNewWork == component.BlockNewWork ||
                 item.BeginShutdown == component.BeginShutdown ||
                 item.Wait == component.Wait)
             {
                 return;
             }
         }
 
         instance.m_components.push_back(component);
     }
 
     void ServerShutdownSynchronization::WaitForShutdown()
     {
         ServerShutdownSynchronization& instance = Instance();
 
         {
             std::lock_guard<std::mutex> lock{ instance.m_threadLock };
             if (!instance.m_shutdownThread.joinable())
             {
                 AICLI_LOG(Core, Warning, << "Attempt to wait for shutdown when shutdown has not been initiated.");
                 return;
             }
         }
 
         instance.m_shutdownComplete.wait();
     }
 
     void ServerShutdownSynchronization::Cancel(CancelReason reason, bool)
     {
         std::lock_guard<std::mutex> lock{ m_threadLock };
 
         if (!m_shutdownThread.joinable())
         {
             m_shutdownThread = std::thread(&ServerShutdownSynchronization::SynchronizeShutdown, this, reason);
         }
     }
 
     ServerShutdownSynchronization::ServerShutdownSynchronization()
     {
         TerminationSignalHandler::Instance()->AddListener(this);
     }
 
     ServerShutdownSynchronization::~ServerShutdownSynchronization()
     {
         TerminationSignalHandler::Instance()->RemoveListener(this);
         if (m_shutdownThread.joinable())
         {
             m_shutdownThread.detach();
         }
     }
 
     ServerShutdownSynchronization& ServerShutdownSynchronization::Instance()
     {
         static ServerShutdownSynchronization s_instance;
         return s_instance;
     }
 
     void ServerShutdownSynchronization::SynchronizeShutdown(CancelReason reason) try
     {
         auto setShutdownComplete = wil::scope_exit([this]() { this->m_shutdownComplete.SetEvent(); });
 
         std::vector<ComponentSystem> components;
         {
             std::lock_guard<std::mutex> lock{ m_componentsLock };
             components = m_components;
         }
 
         for (const auto& component : components)
         {
             if (component.BlockNewWork)
             {
                 component.BlockNewWork(reason);
             }
         }
 
         for (const auto& component : components)
         {
             if (component.BeginShutdown)
             {
                 component.BeginShutdown(reason);
             }
         }
 
         for (const auto& component : components)
         {
             if (component.Wait)
             {
                 component.Wait();
             }
         }
 
         ShutdownCompleteCallback callback = m_callback;
         if (callback)
         {
             callback();
         }
     }
     CATCH_LOG();
 }