Improve COM static store usage (port to 1.11) (#5685)

CP from #5680 

Ported changes backward from the termination handling refactor.

Author: JohnMcPMS

src/AppInstallerCLICore/Commands/TestCommand.cpp

@@ -9,6 +9,7 @@
 #include "TableOutput.h"
 #include "Public/ConfigurationSetProcessorFactoryRemoting.h"
 #include "Workflows/ConfigurationFlow.h"
+#include <winget/RepositorySource.h>
 #include <winrt/Microsoft.Management.Configuration.h>
 
 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<bool (__stdcall *)()>(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<std::unique_ptr<Command>> TestCommand::GetCommands() const
@@ -224,6 +284,7 @@ namespace AppInstaller::CLI
                 std::make_unique<TestAppShutdownCommand>(FullName()),
                 std::make_unique<TestConfigurationExportCommand>(FullName()),
                 std::make_unique<TestConfigurationFindUnitProcessorsCommand>(FullName()),
+                std::make_unique<TestCanUnloadNowCommand>(FullName()),
             });
     }
 

src/AppInstallerCLICore/ExecutionContext.cpp

@@ -7,6 +7,7 @@
 #include "Command.h"
 #include "ExecutionContext.h"
 #include <winget/Checkpoint.h>
+#include <winget/COMStaticStorage.h>
 #include <winget/Reboot.h>
 #include <winget/UserSettings.h>
 #include <winget/NetworkSettings.h>
@@ -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<SignalTerminationHandler> Instance()
             {
-                static SignalTerminationHandler s_instance;
-                return s_instance;
+                struct Singleton : public WinRT::COMStaticStorageBase<SignalTerminationHandler>
+                {
+                    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<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);
+                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
 

src/AppInstallerCLICore/pch.h

@@ -54,4 +54,5 @@
 #pragma warning( pop )
 
 #include <wrl/client.h>
+#include <wrl/implements.h>
 #include <AppxPackaging.h>

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"));
         }
+
+        /// <summary>
+        /// Runs winget test can-unload-now expecting that it cannot be unloaded.
+        /// </summary>
+        [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"));
+        }
     }
 }

src/AppInstallerRepositoryCore/Microsoft/PredefinedInstalledSourceFactory.cpp

@@ -6,7 +6,7 @@
 #include "Microsoft/SQLiteIndex.h"
 #include "Microsoft/SQLiteIndexSource.h"
 #include <winget/ManifestInstaller.h>
-
+#include <winget/COMStaticStorage.h>
 #include <winget/Registry.h>
 #include <AppInstallerArchitecture.h>
 #include <winget/ExperimentalFeature.h>
@@ -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<CachedInstalledIndex>
             {
-                struct Holder : public winrt::implements<Holder, winrt::Windows::Foundation::IInspectable>
-                {
-                    static constexpr std::wstring_view Guid{ L"{48c47064-4fff-4eca-812c-dbb4f33a8fcb}" };
-                    std::shared_ptr<CachedInstalledIndex> m_shared{ std::make_shared<CachedInstalledIndex>() };
-                };
-
-                std::weak_ptr<CachedInstalledIndex> m_weak;
-                winrt::slim_mutex m_lock;
-
-                std::shared_ptr<CachedInstalledIndex> Get()
-                {
-                    {
-                        const std::shared_lock lock{ m_lock };
-                        if (auto cachedIndex = m_weak.lock())
-                        {
-                            return cachedIndex;
-                        }
-                    }
-
-                    auto value = winrt::make_self<Holder>();
-
-                    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<winrt::Windows::Foundation::IInspectable>());
-
-                    m_weak = value->m_shared;
-                    return value->m_shared;
-                }
+                Singleton() : COMStaticStorageBase(L"WindowsPackageManager.CachedInstalledIndex") {}
             };
 
             CachedInstalledIndex()

src/AppInstallerSharedLib/AppInstallerSharedLib.vcxproj

@@ -349,6 +349,7 @@
     <ClInclude Include="Public\winget\AsyncTokens.h" />
     <ClInclude Include="Public\winget\Certificates.h" />
     <ClInclude Include="Public\winget\Compression.h" />
+    <ClInclude Include="Public\winget\COMStaticStorage.h" />
     <ClInclude Include="Public\winget\ConfigurationSetProcessorHandlers.h" />
     <ClInclude Include="Public\winget\Filesystem.h" />
     <ClInclude Include="Public\winget\GroupPolicy.h" />
@@ -379,6 +380,7 @@
     <ClCompile Include="AppInstallerStrings.cpp" />
     <ClCompile Include="Certificates.cpp" />
     <ClCompile Include="Compression.cpp" />
+    <ClCompile Include="COMStaticStorage.cpp" />
     <ClCompile Include="DateTime.cpp" />
     <ClCompile Include="Errors.cpp" />
     <ClCompile Include="Filesystem.cpp" />
@@ -436,4 +438,4 @@
     <Error Condition="!Exists('$(SolutionDir)\packages\Microsoft.Windows.CppWinRT.2.0.230706.1\build\native\Microsoft.Windows.CppWinRT.props')" Text="$([System.String]::Format('$(ErrorText)', '$(SolutionDir)\packages\Microsoft.Windows.CppWinRT.2.0.230706.1\build\native\Microsoft.Windows.CppWinRT.props'))" />
     <Error Condition="!Exists('$(SolutionDir)\packages\Microsoft.Windows.CppWinRT.2.0.230706.1\build\native\Microsoft.Windows.CppWinRT.targets')" Text="$([System.String]::Format('$(ErrorText)', '$(SolutionDir)\packages\Microsoft.Windows.CppWinRT.2.0.230706.1\build\native\Microsoft.Windows.CppWinRT.targets'))" />
   </Target>
-</Project>
+</Project>

src/AppInstallerSharedLib/AppInstallerSharedLib.vcxproj.filters

@@ -140,6 +140,9 @@
     <ClInclude Include="Public\winget\ModuleCountBase.h">
       <Filter>Public\winget</Filter>
     </ClInclude>
+    <ClInclude Include="Public\winget\COMStaticStorage.h">
+      <Filter>Public\winget</Filter>
+    </ClInclude>
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="pch.cpp">
@@ -229,6 +232,9 @@
     <ClCompile Include="SQLiteDynamicStorage.cpp">
       <Filter>SQLite</Filter>
     </ClCompile>
+    <ClCompile Include="COMStaticStorage.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <None Include="PropertySheet.props" />

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<std::wstring> 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();
+}