From 5cb4599f71d804f9ea160615f498e0c900708e58 Mon Sep 17 00:00:00 2001 From: JohnMcPMS Date: Mon, 25 Aug 2025 09:28:18 -0700 Subject: [PATCH] Improve COM static store usage (#5680) The use of the COM static store was causing crashes on long lived processes because the implementation of `DllCanUnloadNow` did not count those objects, nor did it clean them up. This led to our module being unloaded, then when it was reloaded, the recreation of the static store object would invoke the deletion of the old one, which was pointing to the old, unloaded module location. WindowsPackageManager.dll serves many purposes, making the lifetime of statics complex. - It serves as "in-proc" for the CLI - It is the core implementation for the in-proc COM - It is the core implementation for the OOP COM In order to support in-proc COM, we must put static lifetime COM objects in the static store. But in order to support unloading, we must also clean them up. Additionally, we don't want to claim to be in use if the only active objects are our statics (which are typically just event handlers and their owners). We already use the WRL object count to track OOP COM server lifetime, and similarly we use it to implement `DllCanUnloadNow`. This is the count externally owned objects; those that the client has requested directly or indirectly. The major change is to remove all of our static store objects when WRL says we have no more externally owned. This is achieved by tracking the names of the objects that we insert and attempting to remove them when appropriate. The original change to use the static store was templatized and reused to hold the termination signal handler. The new test uses the CLI to validate that the implementation for `DllCanUnloadNow` (`WindowsPackageManagerInProcModuleTerminate`) detects the unload state and properly destroys the relevant objects. This is done by checking that there are internal objects allocated before the call, but none after. --- .../Commands/TestCommand.cpp | 61 ++++++++++++++ src/AppInstallerCLICore/ExecutionContext.cpp | 40 +++++---- src/AppInstallerCLICore/pch.h | 1 + .../AppShutdownTests.cs | 26 ++++++ .../PredefinedInstalledSourceFactory.cpp | 38 +-------- .../AppInstallerSharedLib.vcxproj | 4 +- .../AppInstallerSharedLib.vcxproj.filters | 6 ++ .../COMStaticStorage.cpp | 38 +++++++++ .../Public/winget/COMStaticStorage.h | 83 +++++++++++++++++++ src/AppInstallerSharedLib/pch.h | 4 +- src/WindowsPackageManager/main.cpp | 20 +++-- 11 files changed, 265 insertions(+), 56 deletions(-) create mode 100644 src/AppInstallerSharedLib/COMStaticStorage.cpp create mode 100644 src/AppInstallerSharedLib/Public/winget/COMStaticStorage.h diff --git a/src/AppInstallerCLICore/Commands/TestCommand.cpp b/src/AppInstallerCLICore/Commands/TestCommand.cpp index 7dd0a9877f..fa8a130233 100644 --- a/src/AppInstallerCLICore/Commands/TestCommand.cpp +++ b/src/AppInstallerCLICore/Commands/TestCommand.cpp @@ -9,6 +9,7 @@ #include "TableOutput.h" #include "Public/ConfigurationSetProcessorFactoryRemoting.h" #include "Workflows/ConfigurationFlow.h" +#include #include using namespace AppInstaller::CLI::Workflow; @@ -215,6 +216,65 @@ namespace AppInstaller::CLI InvokeFindUnitProcessors; } }; + + struct TestCanUnloadNowCommand final : public Command + { + TestCanUnloadNowCommand(std::string_view parent) : Command("can-unload-now", {}, parent, Visibility::Hidden) {} + + Resource::LocString ShortDescription() const override + { + return "Test DllCanUnloadNow"_lis; + } + + Resource::LocString LongDescription() const override + { + return "Verifies that the function that implements the inproc DllCanUnloadNow properly blocks unload due to static storage object."_lis; + } + + protected: + void ExecuteInternal(Execution::Context& context) const override + { + Repository::Source source{ Repository::PredefinedSource::Installed }; + + ProgressCallback progress; + source.Open(progress); + + HMODULE self = GetModuleHandle(L"WindowsPackageManager.dll"); + if (!self) + { + LogAndReport(context, "Couldn't get WindowsPackageManager module"); + return; + } + + auto WindowsPackageManagerInProcModuleTerminate = reinterpret_cast(GetProcAddress(self, "WindowsPackageManagerInProcModuleTerminate")); + + // Report the object counts, attempt to terminate, report the object counts again + ReportObjectCounts(context); + LogAndReport(context, WindowsPackageManagerInProcModuleTerminate() ? "DllCanUnloadNow" : "DllCannotUnloadNow"); + ReportObjectCounts(context); + } + + private: + void ReportObjectCounts(Execution::Context& context) const + { + std::ostringstream stream; + stream << "Internal objects: " << GetInternalObjectCount() << '\n'; + stream << "External objects: " << GetExternalObjectCount(); + + LogAndReport(context, stream.str()); + } + + uint32_t GetInternalObjectCount() const + { + return winrt::get_module_lock().operator unsigned int(); + } + + unsigned long GetExternalObjectCount() const + { + auto module = Microsoft::WRL::GetModuleBase(); + return module ? module->GetObjectCount() : 0; + } + }; } std::vector> TestCommand::GetCommands() const @@ -224,6 +284,7 @@ namespace AppInstaller::CLI std::make_unique(FullName()), std::make_unique(FullName()), std::make_unique(FullName()), + std::make_unique(FullName()), }); } diff --git a/src/AppInstallerCLICore/ExecutionContext.cpp b/src/AppInstallerCLICore/ExecutionContext.cpp index a3eaea509c..803372cd5d 100644 --- a/src/AppInstallerCLICore/ExecutionContext.cpp +++ b/src/AppInstallerCLICore/ExecutionContext.cpp @@ -7,6 +7,7 @@ #include "Command.h" #include "ExecutionContext.h" #include +#include #include #include #include @@ -22,10 +23,15 @@ namespace AppInstaller::CLI::Execution // Type to contain the CTRL signal and window messages handler. struct SignalTerminationHandler { - static SignalTerminationHandler& Instance() + static std::shared_ptr Instance() { - static SignalTerminationHandler s_instance; - return s_instance; + struct Singleton : public WinRT::COMStaticStorageBase + { + Singleton() : COMStaticStorageBase(L"WindowsPackageManager.SignalTerminationHandler") {} + }; + + static Singleton s_instance; + return s_instance.Get(); } void AddContext(Context* context) @@ -42,8 +48,14 @@ namespace AppInstaller::CLI::Execution std::lock_guard 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); + if (itr == m_contexts.end()) + { + AICLI_LOG(CLI, Warning, << "SignalTerminationHandler::RemoveContext did not find requested object"); + } + else + { + m_contexts.erase(itr); + } } void StartAppShutdown() @@ -66,7 +78,6 @@ namespace AppInstaller::CLI::Execution } #endif - private: SignalTerminationHandler() { if (Runtime::IsRunningAsAdmin() && Runtime::IsRunningInPackagedContext()) @@ -81,7 +92,7 @@ namespace AppInstaller::CLI::Execution auto progress = args.Progress(); if (progress > minProgress) { - SignalTerminationHandler::Instance().StartAppShutdown(); + SignalTerminationHandler::Instance()->StartAppShutdown(); } }); } @@ -110,13 +121,14 @@ namespace AppInstaller::CLI::Execution // if there's no call to join. if (m_windowThread.joinable()) { - m_windowThread.join(); + m_windowThread.detach(); } } + private: static BOOL WINAPI StaticCtrlHandlerFunction(DWORD ctrlType) { - return Instance().CtrlHandlerFunction(ctrlType); + return Instance()->CtrlHandlerFunction(ctrlType); } static LRESULT WINAPI WindowMessageProcedure(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam) @@ -125,7 +137,7 @@ namespace AppInstaller::CLI::Execution switch (uMsg) { case WM_QUERYENDSESSION: - SignalTerminationHandler::Instance().StartAppShutdown(); + SignalTerminationHandler::Instance()->StartAppShutdown(); return TRUE; case WM_ENDSESSION: case WM_CLOSE: @@ -267,11 +279,11 @@ namespace AppInstaller::CLI::Execution if (add) { - SignalTerminationHandler::Instance().AddContext(context); + SignalTerminationHandler::Instance()->AddContext(context); } else { - SignalTerminationHandler::Instance().RemoveContext(context); + SignalTerminationHandler::Instance()->RemoveContext(context); } } @@ -490,12 +502,12 @@ namespace AppInstaller::CLI::Execution HWND GetWindowHandle() { - return SignalTerminationHandler::Instance().GetWindowHandle(); + return SignalTerminationHandler::Instance()->GetWindowHandle(); } bool WaitForAppShutdownEvent() { - return SignalTerminationHandler::Instance().WaitForAppShutdownEvent(); + return SignalTerminationHandler::Instance()->WaitForAppShutdownEvent(); } #endif diff --git a/src/AppInstallerCLICore/pch.h b/src/AppInstallerCLICore/pch.h index e11cf29652..94e30d058d 100644 --- a/src/AppInstallerCLICore/pch.h +++ b/src/AppInstallerCLICore/pch.h @@ -54,4 +54,5 @@ #pragma warning( pop ) #include +#include #include diff --git a/src/AppInstallerCLIE2ETests/AppShutdownTests.cs b/src/AppInstallerCLIE2ETests/AppShutdownTests.cs index 4534282a83..49d9727cf9 100644 --- a/src/AppInstallerCLIE2ETests/AppShutdownTests.cs +++ b/src/AppInstallerCLIE2ETests/AppShutdownTests.cs @@ -7,6 +7,7 @@ namespace AppInstallerCLIE2ETests { using System; + using System.Diagnostics; using System.IO; using System.Threading; using System.Threading.Tasks; @@ -35,6 +36,12 @@ public void RegisterApplicationTest() Assert.Ignore("This test won't work on Window Server as non-admin"); } + if (!Environment.Is64BitProcess) + { + // My guess is that HAM terminates us faster after the CTRL-C on x86... + Assert.Ignore("This test is flaky when run as x86."); + } + if (string.IsNullOrEmpty(TestSetup.Parameters.AICLIPackagePath)) { throw new NullReferenceException("AICLIPackagePath"); @@ -95,5 +102,24 @@ public void RegisterApplicationTest() // Look for the output. Assert.True(testCmdTask.Result.StdOut.Contains("Succeeded waiting for app shutdown event")); } + + /// + /// Runs winget test can-unload-now expecting that it cannot be unloaded. + /// + [Test] + public void CanUnloadNowTest() + { + var result = TestCommon.RunAICLICommand("test", "can-unload-now --verbose"); + + var lines = result.StdOut.Split('\n', StringSplitOptions.TrimEntries | StringSplitOptions.RemoveEmptyEntries); + + Assert.AreEqual(5, lines.Length); + Assert.True(lines[0].Contains("Internal objects:")); + Assert.False(lines[0].Contains("Internal objects: 0")); + Assert.True(lines[1].Contains("External objects: 0")); + Assert.True(lines[2].Contains("DllCanUnloadNow")); + Assert.True(lines[3].Contains("Internal objects: 0")); + Assert.True(lines[4].Contains("External objects: 0")); + } } } \ No newline at end of file diff --git a/src/AppInstallerRepositoryCore/Microsoft/PredefinedInstalledSourceFactory.cpp b/src/AppInstallerRepositoryCore/Microsoft/PredefinedInstalledSourceFactory.cpp index 38a15f003a..1d1b559f45 100644 --- a/src/AppInstallerRepositoryCore/Microsoft/PredefinedInstalledSourceFactory.cpp +++ b/src/AppInstallerRepositoryCore/Microsoft/PredefinedInstalledSourceFactory.cpp @@ -6,7 +6,7 @@ #include "Microsoft/SQLiteIndex.h" #include "Microsoft/SQLiteIndexSource.h" #include - +#include #include #include #include @@ -262,41 +262,9 @@ namespace AppInstaller::Repository::Microsoft struct CachedInstalledIndex { - // https://devblogs.microsoft.com/oldnewthing/20210215-00/?p=104865 - struct Singleton + struct Singleton : public WinRT::COMStaticStorageBase { - struct Holder : public winrt::implements - { - static constexpr std::wstring_view Guid{ L"{48c47064-4fff-4eca-812c-dbb4f33a8fcb}" }; - std::shared_ptr m_shared{ std::make_shared() }; - }; - - std::weak_ptr m_weak; - winrt::slim_mutex m_lock; - - std::shared_ptr Get() - { - { - const std::shared_lock lock{ m_lock }; - if (auto cachedIndex = m_weak.lock()) - { - return cachedIndex; - } - } - - auto value = winrt::make_self(); - - const std::shared_lock lock{ m_lock }; - if (auto cachedIndex = m_weak.lock()) - { - return cachedIndex; - } - - winrt::Windows::ApplicationModel::Core::CoreApplication::Properties().Insert(Holder::Guid, value.as()); - - m_weak = value->m_shared; - return value->m_shared; - } + Singleton() : COMStaticStorageBase(L"WindowsPackageManager.CachedInstalledIndex") {} }; CachedInstalledIndex() diff --git a/src/AppInstallerSharedLib/AppInstallerSharedLib.vcxproj b/src/AppInstallerSharedLib/AppInstallerSharedLib.vcxproj index 05603d1007..ea81c95fe6 100644 --- a/src/AppInstallerSharedLib/AppInstallerSharedLib.vcxproj +++ b/src/AppInstallerSharedLib/AppInstallerSharedLib.vcxproj @@ -349,6 +349,7 @@ + @@ -379,6 +380,7 @@ + @@ -436,4 +438,4 @@ - + \ No newline at end of file diff --git a/src/AppInstallerSharedLib/AppInstallerSharedLib.vcxproj.filters b/src/AppInstallerSharedLib/AppInstallerSharedLib.vcxproj.filters index dce8ddd078..31ad082c5a 100644 --- a/src/AppInstallerSharedLib/AppInstallerSharedLib.vcxproj.filters +++ b/src/AppInstallerSharedLib/AppInstallerSharedLib.vcxproj.filters @@ -140,6 +140,9 @@ Public\winget + + Public\winget + @@ -229,6 +232,9 @@ SQLite + + Source Files + diff --git a/src/AppInstallerSharedLib/COMStaticStorage.cpp b/src/AppInstallerSharedLib/COMStaticStorage.cpp new file mode 100644 index 0000000000..f271ab1823 --- /dev/null +++ b/src/AppInstallerSharedLib/COMStaticStorage.cpp @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +#include "pch.h" +#include "Public/winget/COMStaticStorage.h" + +namespace AppInstaller::WinRT +{ + COMStaticStorageStatics& COMStaticStorageStatics::Instance() + { + static COMStaticStorageStatics s_instance; + return s_instance; + } + + void COMStaticStorageStatics::AddStaticStorageItem(const winrt::hstring& name, const winrt::Windows::Foundation::IInspectable& item) + { + COMStaticStorageStatics& instance = Instance(); + const winrt::slim_lock_guard lock{ instance.m_lock }; + winrt::Windows::ApplicationModel::Core::CoreApplication::Properties().Insert(name, item); + instance.m_items.emplace(std::wstring{ name }); + } + + void COMStaticStorageStatics::ResetAll() try + { + COMStaticStorageStatics& instance = Instance(); + std::set localItems; + + { + const winrt::slim_lock_guard lock{ instance.m_lock }; + instance.m_items.swap(localItems); + } + + for (const auto& item : localItems) + { + winrt::Windows::ApplicationModel::Core::CoreApplication::Properties().TryRemove(item); + } + } + CATCH_LOG(); +} diff --git a/src/AppInstallerSharedLib/Public/winget/COMStaticStorage.h b/src/AppInstallerSharedLib/Public/winget/COMStaticStorage.h new file mode 100644 index 0000000000..a7ab1d03a7 --- /dev/null +++ b/src/AppInstallerSharedLib/Public/winget/COMStaticStorage.h @@ -0,0 +1,83 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +#pragma once +#include +#include +#include +#include +#include +#include + +namespace AppInstaller::WinRT +{ + // Contains registration for static storage so that they can be cleared. + struct COMStaticStorageStatics + { + // Adds a static storage key to the set of known items. + static void AddStaticStorageItem(const winrt::hstring& name, const winrt::Windows::Foundation::IInspectable& item); + + // Removes all known static storage items. + static void ResetAll(); + + private: + COMStaticStorageStatics() = default; + + static COMStaticStorageStatics& Instance(); + + winrt::slim_mutex m_lock; + std::set m_items; + }; + + // https://devblogs.microsoft.com/oldnewthing/20210215-00/?p=104865 + // Base class for an object that needs to live in the COM static store. + // An object needs to use this if it has: + // - static lifetime + // - references to externally implemented COM objects + // + // Additionally, it should *not* contain references to WRL counted objects implemented by this module. + // If it does, it will prevent the module from being unloaded until COM is uninitialized, which is often never. + template + struct COMStaticStorageBase + { + private: + struct DataHolder : public winrt::implements + { + std::shared_ptr m_shared{ std::make_shared() }; + }; + + std::weak_ptr m_weak; + winrt::slim_mutex m_lock; + winrt::hstring m_name; + + public: + COMStaticStorageBase(std::wstring_view name) : m_name(name) {} + + std::shared_ptr Get() + { + { + const std::shared_lock lock{ m_lock }; + if (auto cached = m_weak.lock()) + { + return cached; + } + } + + auto value = winrt::make_self(); + + const winrt::slim_lock_guard lock{ m_lock }; + if (auto cached = m_weak.lock()) + { + return cached; + } + + COMStaticStorageStatics::AddStaticStorageItem(m_name, value.as()); + m_weak = value->m_shared; + return value->m_shared; + } + + void Reset() + { + winrt::Windows::ApplicationModel::Core::CoreApplication::Properties().TryRemove(m_name); + } + }; +} diff --git a/src/AppInstallerSharedLib/pch.h b/src/AppInstallerSharedLib/pch.h index fd3165082e..22186ce6ba 100644 --- a/src/AppInstallerSharedLib/pch.h +++ b/src/AppInstallerSharedLib/pch.h @@ -58,8 +58,10 @@ #include #pragma warning( pop ) -#include +#include +#include #include #include +#include #include #include diff --git a/src/WindowsPackageManager/main.cpp b/src/WindowsPackageManager/main.cpp index 751d75caa2..1612aeaac1 100644 --- a/src/WindowsPackageManager/main.cpp +++ b/src/WindowsPackageManager/main.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include using namespace winrt::Microsoft::Management::Deployment; @@ -98,12 +99,21 @@ extern "C" { try { - return ::Microsoft::WRL::Module<::Microsoft::WRL::ModuleType::InProc>::GetModule().Terminate(); - } - catch (...) - { - return false; + // The WRL object count is used to track externally visible objects, which largely means objects created with the `wil::details::module_count_wrapper` type wrapper. + // Configuration objects use a composition based tracking that is similar in nature (only when OOP). + // + // In-proc DllCanUnloadNow should not be blocked by our internal objects, but they must be destroyed on unload or a future reload will attempt to destroy them + // and our module may have moved. So when we don't have any more objects that we gave to callers, remove all of our static lifetime objects and indicate + // that we can now be unloaded. + if (::Microsoft::WRL::Module<::Microsoft::WRL::ModuleType::InProc>::GetModule().Terminate()) + { + AppInstaller::WinRT::COMStaticStorageStatics::ResetAll(); + return true; + } } + catch (...) {} + + return false; } WINDOWS_PACKAGE_MANAGER_API WindowsPackageManagerInProcModuleGetClassObject(