winget-cli

ContextOrchestrator.cpp (23715B)

---

 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 #include "pch.h"
 #include "ExecutionContext.h"
 #include "ContextOrchestrator.h"
 #include "COMContext.h"
 #include "Commands/COMCommand.h"
 #include "Public/ShutdownMonitoring.h"
 #include "winget/UserSettings.h"
 #include <Commands/RootCommand.h>
 
 namespace AppInstaller::CLI::Execution
 {
     namespace
     {
         // Operation command queue used by install, uninstall and repair commands.
         constexpr static std::string_view OperationCommandQueueName = "operation"sv;
 
         // Callback function used by worker threads in the queue.
         // context must be a pointer to a queue item.
         void CALLBACK OrchestratorQueueWorkCallback(PTP_CALLBACK_INSTANCE, PVOID context, PTP_WORK)
         {
             auto queueItem = reinterpret_cast<OrchestratorQueueItem*>(context);
             auto queue = queueItem->GetCurrentQueue();
             if (queue)
             {
                 queue->RunItem(queueItem->GetId());
             }
         }
 
         // Get command queue name based on command name.
         std::string_view GetCommandQueueName(std::string_view commandName)
         {
             if (commandName == COMInstallCommand::CommandName || commandName == COMUninstallCommand::CommandName || commandName == COMRepairCommand::CommandName)
             {
                 return OperationCommandQueueName;
             }
 
             return commandName;
         }
     }
 
     ContextOrchestrator& ContextOrchestrator::Instance()
     {
         static ContextOrchestrator s_instance;
         return s_instance;
     }
 
     ContextOrchestrator::ContextOrchestrator() : ContextOrchestrator(std::thread::hardware_concurrency()) {}
 
     ContextOrchestrator::ContextOrchestrator(unsigned int hardwareConcurrency)
     {
         ProgressCallback progress;
         m_installingWriteableSource = Repository::Source(Repository::PredefinedSource::Installing);
         m_installingWriteableSource.Open(progress);
 
         // Decide how many threads to use for each command.
         // We always allow only one install at a time.
         // For download, if we can find the number of supported concurrent threads,
         // use that as the maximum (up to 3); otherwise use a single thread.
         const UINT32 maxDownloadThreads = 3;
         const UINT32 operationThreads = 1;
         const UINT32 downloadThreads = std::min(hardwareConcurrency > 1 ? hardwareConcurrency - 1 : 1, maxDownloadThreads);
 
         AddCommandQueue(COMDownloadCommand::CommandName, downloadThreads);
         AddCommandQueue(OperationCommandQueueName, operationThreads);
     }
 
     void ContextOrchestrator::AddCommandQueue(std::string_view commandName, UINT32 allowedThreads)
     {
         std::lock_guard<std::mutex> lockQueue{ m_queueLock };
         m_commandQueues.emplace(commandName, std::make_unique<OrchestratorQueue>(*this, commandName, allowedThreads));
     }
 
     _Requires_lock_held_(m_queueLock)
     std::shared_ptr<OrchestratorQueueItem> ContextOrchestrator::FindById(const OrchestratorQueueItemId& comparisonQueueItemId)
     {
         for (const auto& queue : m_commandQueues)
         {
             auto item = queue.second->FindById(comparisonQueueItemId);
             if (item)
             {
                 return item;
             }
         }
 
         return {};
     }
 
     void ContextOrchestrator::EnqueueAndRunItem(const std::shared_ptr<OrchestratorQueueItem>& item)
     {
         std::lock_guard<std::mutex> lockQueue{ m_queueLock };
 
         if (item->IsOnFirstCommand())
         {
             // Directly error on attempting to enqueue first time
             THROW_HR_IF(ToHRESULT(m_disabledReason), !m_enabled);
 
             THROW_HR_IF(HRESULT_FROM_WIN32(ERROR_INSTALL_ALREADY_RUNNING), FindById(item->GetId()));
 
             // Log the beginning of the item
             item->GetContext().GetThreadGlobals().GetTelemetryLogger().LogCommand(item->GetItemCommandName());
         }
         else if (!m_enabled)
         {
             // On subsequent command enqueues, cancel and complete the item
             item->GetContext().Cancel(m_disabledReason, true);
             item->HandleItemCompletion(*this);
         }
 
         std::string commandQueueName{ GetCommandQueueName(item->GetNextCommand().Name()) };
         m_commandQueues.at(commandQueueName)->EnqueueAndRunItem(item);
     }
 
     void ContextOrchestrator::RemoveItemInState(const OrchestratorQueueItem& item, OrchestratorQueueItemState state)
     {
         std::lock_guard<std::mutex> lockQueue{ m_queueLock };
         for (const auto& queue : m_commandQueues)
         {
             if (queue.second->RemoveItemInState(item, state, true))
             {
                 return;
             }
         }
     }
 
     void ContextOrchestrator::CancelQueueItem(const OrchestratorQueueItem& item)
     {
         // Always cancel the item, even if it isn't running yet, to get the terminationHR set correctly.
         item.GetContext().Cancel(CancelReason::Abort, true);
 
         RemoveItemInState(item, OrchestratorQueueItemState::Queued);
     }
 
     std::shared_ptr<OrchestratorQueueItem> ContextOrchestrator::GetQueueItem(const OrchestratorQueueItemId& queueItemId)
     {
         std::lock_guard<std::mutex> lock{ m_queueLock };
 
         return FindById(queueItemId);
     }
 
     void ContextOrchestrator::AddItemManifestToInstallingSource(const OrchestratorQueueItem& queueItem)
     {
         if (queueItem.IsApplicableForInstallingSource())
         {
             const auto& manifest = queueItem.GetContext().Get<Execution::Data::Manifest>();
             m_installingWriteableSource.AddPackageVersion(manifest, std::filesystem::path{ manifest.Id + '.' + manifest.Version });
         }
     }
 
     void ContextOrchestrator::RemoveItemManifestFromInstallingSource(const OrchestratorQueueItem& queueItem)
     {
         if (queueItem.IsApplicableForInstallingSource())
         {
             const auto& manifest = queueItem.GetContext().Get<Execution::Data::Manifest>();
             m_installingWriteableSource.RemovePackageVersion(manifest, std::filesystem::path{ manifest.Id + '.' + manifest.Version });
         }
     }
 
     void ContextOrchestrator::RegisterForShutdownSynchronization()
     {
         static std::once_flag registerComponentOnceFlag;
         std::call_once(registerComponentOnceFlag,
             [&]()
             {
                 using namespace ShutdownMonitoring;
 
                 ServerShutdownSynchronization::ComponentSystem component;
                 component.BlockNewWork = StaticDisable;
                 component.BeginShutdown = StaticCancelQueuedItems;
                 component.Wait = StaticWaitForRunningItems;
 
                 ServerShutdownSynchronization::AddComponent(component);
             });
     }
 
     void ContextOrchestrator::StaticDisable(CancelReason reason)
     {
         Instance().Disable(reason);
     }
 
     void ContextOrchestrator::StaticCancelQueuedItems(CancelReason reason)
     {
         Instance().CancelQueuedItems(reason);
     }
 
     void ContextOrchestrator::StaticWaitForRunningItems()
     {
         Instance().WaitForRunningItems();
     }
 
     void ContextOrchestrator::Disable(CancelReason reason)
     {
         std::lock_guard<std::mutex> lock{ m_queueLock };
         m_enabled = false;
         m_disabledReason = reason;
     }
 
     void ContextOrchestrator::CancelQueuedItems(CancelReason reason)
     {
         std::lock_guard<std::mutex> lock{ m_queueLock };
         for (const auto& queue : m_commandQueues)
         {
             queue.second->CancelAllItems(reason);
         }
     }
 
     void ContextOrchestrator::WaitForRunningItems()
     {
         std::lock_guard<std::mutex> lock{ m_queueLock };
         for (const auto& queue : m_commandQueues)
         {
             queue.second->WaitForEmptyQueue();
         }
     }
 
     bool ContextOrchestrator::WaitForRunningItems(DWORD timeoutMilliseconds)
     {
         std::lock_guard<std::mutex> lock{ m_queueLock };
         for (const auto& queue : m_commandQueues)
         {
             if (!queue.second->WaitForEmptyQueue(timeoutMilliseconds))
             {
                 return false;
             }
         }
 
         return true;
     }
 
     std::string ContextOrchestrator::GetStatusString()
     {
         std::ostringstream stream;
 
         std::lock_guard<std::mutex> lock{ m_queueLock };
 
         if (!m_enabled)
         {
             stream << "Disabled due to " << ToIntegral(m_disabledReason) << std::endl;
         }
 
         for (const auto& queue : m_commandQueues)
         {
             stream << queue.second->GetStatusString();
         }
 
         return stream.str();
     }
 
     _Requires_lock_held_(m_itemLock)
     std::deque<std::shared_ptr<OrchestratorQueueItem>>::iterator OrchestratorQueue::FindIteratorById(const OrchestratorQueueItemId& comparisonQueueItemId)
     {
         return std::find_if(m_queueItems.begin(), m_queueItems.end(), [&comparisonQueueItemId](const std::shared_ptr<OrchestratorQueueItem>& item) {return (item->GetId().IsSame(comparisonQueueItemId)); });
     }
 
     _Requires_lock_held_(m_itemLock)
     std::shared_ptr<OrchestratorQueueItem> OrchestratorQueue::FindById(const OrchestratorQueueItemId& comparisonQueueItemId)
     {
         auto itr = FindIteratorById(comparisonQueueItemId);
         if (itr != m_queueItems.end())
         {
             return *itr;
         }
 
         return {};
     }
 
     void OrchestratorQueue::EnqueueItem(const std::shared_ptr<OrchestratorQueueItem>& item)
     {
         {
             std::lock_guard<std::mutex> lockQueue{ m_itemLock };
             m_queueItems.push_back(item);
             m_queueEmpty.ResetEvent();
         }
 
         // Add the package to the Installing source so that it can be queried using the Source interface.
         // Only do this the first time the item is queued.
         if (item->IsOnFirstCommand())
         {
             try
             {
                 m_orchestrator.AddItemManifestToInstallingSource(*item);
             }
             catch (...)
             {
                 std::lock_guard<std::mutex> lockQueue{ m_itemLock };
                 auto itr = FindIteratorById(item->GetId());
                 if (itr != m_queueItems.end())
                 {
                     m_queueItems.erase(itr);
 
                     if (m_queueItems.empty())
                     {
                         m_queueEmpty.SetEvent();
                     }
                 }
                 throw;
             }
         }
 
         {
             std::lock_guard<std::mutex> lockQueue{ m_itemLock };
             item->SetState(OrchestratorQueueItemState::Queued);
         }
     }
 
     OrchestratorQueue::OrchestratorQueue(ContextOrchestrator& orchestrator, std::string_view commandName, UINT32 allowedThreads) :
         m_orchestrator(orchestrator), m_commandName(commandName), m_allowedThreads(allowedThreads)
     {
         m_threadPool.reset(CreateThreadpool(nullptr));
         THROW_LAST_ERROR_IF_NULL(m_threadPool);
         m_threadPoolCleanupGroup.reset(CreateThreadpoolCleanupGroup());
         THROW_LAST_ERROR_IF_NULL(m_threadPoolCleanupGroup);
         InitializeThreadpoolEnvironment(&m_threadPoolCallbackEnviron);
         SetThreadpoolCallbackPool(&m_threadPoolCallbackEnviron, m_threadPool.get());
         SetThreadpoolCallbackCleanupGroup(&m_threadPoolCallbackEnviron, m_threadPoolCleanupGroup.get(), nullptr);
 
         SetThreadpoolThreadMaximum(m_threadPool.get(), m_allowedThreads);
         THROW_LAST_ERROR_IF(!SetThreadpoolThreadMinimum(m_threadPool.get(), 1));
     }
 
     OrchestratorQueue::~OrchestratorQueue()
     {
         CloseThreadpoolCleanupGroupMembers(m_threadPoolCleanupGroup.get(), false, nullptr);
     }
 
     void OrchestratorQueue::EnqueueAndRunItem(const std::shared_ptr<OrchestratorQueueItem>& item)
     {
         EnqueueItem(item);
 
         item->SetCurrentQueue(this);
         auto work = CreateThreadpoolWork(OrchestratorQueueWorkCallback, item.get(), &m_threadPoolCallbackEnviron);
         SubmitThreadpoolWork(work);
     }
 
     void OrchestratorQueue::RunItem(const OrchestratorQueueItemId& itemId)
     {
         try
         {
             std::shared_ptr<OrchestratorQueueItem> item;
             bool isCancelled = false;
 
             // Try to find the item in the queue.
             {
                 std::lock_guard<std::mutex> lockQueue{ m_itemLock };
                 item = FindById(itemId);
 
                 if (!item)
                 {
                     // Item should be in the queue; this shouldn't happen.
                     return;
                 }
 
                 // Only run if the item is queued and not cancelled.
                 if (item->GetState() == OrchestratorQueueItemState::Queued)
                 {
                     // Mark it as running so that it cannot be cancelled by other threads.
                     item->SetState(OrchestratorQueueItemState::Running);
                 }
                 else if (item->GetState() == OrchestratorQueueItemState::Cancelled)
                 {
                     isCancelled = true;
                 }
             }
 
             if (isCancelled)
             {
                 // Do this separate from above block as the Remove function needs to manage the lock.
                 RemoveItemInState(*item, OrchestratorQueueItemState::Cancelled, true);
             }
 
             // Get the item's command and execute it.
             HRESULT exceptionHR = S_OK;
             try
             {
                 std::unique_ptr<Command> command = item->PopNextCommand();
 
                 std::unique_ptr<AppInstaller::ThreadLocalStorage::PreviousThreadGlobals> setThreadGlobalsToPreviousState = item->GetContext().SetForCurrentThread();
 
                 command->ValidateArguments(item->GetContext().Args);
 
                 item->GetContext().EnableSignalTerminationHandler();
 
                 ::AppInstaller::CLI::ExecuteWithoutLoggingSuccess(item->GetContext(), command.get());
             }
             WINGET_CATCH_STORE(exceptionHR, APPINSTALLER_CLI_ERROR_COMMAND_FAILED);
 
             if (FAILED(exceptionHR))
             {
                 // Set the termination hr directly from any exception that escaped so that the context always 
                 // has the result of the operation no matter how it failed.
                 item->GetContext().SetTerminationHR(exceptionHR);
             }
 
             item->GetContext().EnableSignalTerminationHandler(false);
 
             if (FAILED(item->GetContext().GetTerminationHR()) || item->IsComplete())
             {
                 if (SUCCEEDED(item->GetContext().GetTerminationHR()))
                 {
                     item->GetContext().GetThreadGlobals().GetTelemetryLogger().LogCommandSuccess(item->GetItemCommandName());
                 }
 
                 RemoveItemInState(*item, OrchestratorQueueItemState::Running, true);
             }
             else
             {
                 // Remove item from this queue and add it to the queue for the next command.
                 RemoveItemInState(*item, OrchestratorQueueItemState::Running, false);
                 m_orchestrator.EnqueueAndRunItem(item);
             }
         }
         catch (...)
         {
         }
     }
 
     void OrchestratorQueue::CancelAllItems(CancelReason reason)
     {
         std::lock_guard<std::mutex> lockQueue{ m_itemLock };
 
         for (auto itr = m_queueItems.begin(); itr != m_queueItems.end(); itr++)
         {
             auto& item = *itr;
 
             item->GetContext().Cancel(reason, true);
 
             // This mimics ContextOrchestrator::CancelQueueItem, which speeds up the process of cancelling queued items
             if (item->GetState() == OrchestratorQueueItemState::Queued)
             {
                 item->SetState(OrchestratorQueueItemState::Cancelled);
                 item->HandleItemCompletion(m_orchestrator);
             }
         }
     }
 
     void OrchestratorQueue::WaitForEmptyQueue()
     {
         m_queueEmpty.wait();
     }
 
     bool OrchestratorQueue::WaitForEmptyQueue(DWORD timeoutMilliseconds)
     {
         return m_queueEmpty.wait(timeoutMilliseconds);
     }
 
     std::string OrchestratorQueue::GetStatusString()
     {
         std::ostringstream stream;
         stream << m_commandName << '[' << m_allowedThreads << "]\n";
 
         std::map<OrchestratorQueueItemState, size_t> stateCounts;
         stateCounts[OrchestratorQueueItemState::NotQueued] = 0;
         stateCounts[OrchestratorQueueItemState::Queued] = 0;
         stateCounts[OrchestratorQueueItemState::Running] = 0;
         stateCounts[OrchestratorQueueItemState::Cancelled] = 0;
 
         {
             std::lock_guard<std::mutex> lock{ m_itemLock };
 
             for (const auto& item : m_queueItems)
             {
                 stateCounts[item->GetState()] += 1;
             }
         }
 
         for (const auto& stateCount : stateCounts)
         {
             stream << "  " << ToString(stateCount.first) << " : " << stateCount.second << std::endl;
         }
 
         return stream.str();
     }
 
     bool OrchestratorQueue::RemoveItemInState(const OrchestratorQueueItem& item, OrchestratorQueueItemState state, bool isGlobalRemove)
     {
         // OrchestratorQueueItemState::Running items should only be removed by the thread that ran the item.
         // Queued items can be removed by any thread.
         // NotQueued items should not be removed since, if found in the queue, they are in the process of being queued by another thread.
         bool foundItem = false;
 
         {
             std::lock_guard<std::mutex> lockQueue{ m_itemLock };
 
             // Look for the item. It's ok if the item is not found since multiple listeners may try to remove the same item.
             auto itr = FindIteratorById(item.GetId());
             if (itr != m_queueItems.end() && (*itr)->GetState() == state)
             {
                 foundItem = true;
 
                 // The item must only be removed from the queue by the thread that runs
                 // it, because the callback uses it. If any other thread tries to remove
                 // it, we simply mark it as cancelled.
                 if (state == OrchestratorQueueItemState::Running || state == OrchestratorQueueItemState::Cancelled)
                 {
                     (*itr)->SetCurrentQueue(nullptr);
                     m_queueItems.erase(itr);
 
                     if (m_queueItems.empty())
                     {
                         m_queueEmpty.SetEvent();
                     }
                 }
                 else if (state == OrchestratorQueueItemState::Queued)
                 {
                     (*itr)->SetState(OrchestratorQueueItemState::Cancelled);
                 }
             }
         }
 
         if (foundItem && isGlobalRemove)
         {
             item.HandleItemCompletion(m_orchestrator);
         }
 
         return foundItem;
     }
 
     bool OrchestratorQueueItemId::IsSame(const OrchestratorQueueItemId& comparedId) const
     {
         return ((GetPackageId() == comparedId.GetPackageId()) && 
                 (GetSourceId() == comparedId.GetSourceId()));
     }
 
     std::string_view OrchestratorQueueItem::GetItemCommandName() const
     {
         // The goal is that these should match the winget.exe commands for easy correlation.
         switch (m_operationType)
         {
         case PackageOperationType::Search: return "root:search"sv;
         case PackageOperationType::Install: return "root:install"sv;
         case PackageOperationType::Upgrade: return "root:upgrade"sv;
         case PackageOperationType::Uninstall: return "root:uninstall"sv;
         case PackageOperationType::Download: return "root:download"sv;
         case PackageOperationType::Repair: return "root:repair"sv;
         default: return "unknown";
         }
     }
 
     void OrchestratorQueueItem::HandleItemCompletion(ContextOrchestrator& orchestrator) const
     {
         orchestrator.RemoveItemManifestFromInstallingSource(*this);
         GetCompletedEvent().SetEvent();
     }
 
     std::unique_ptr<OrchestratorQueueItem> OrchestratorQueueItemFactory::CreateItemForInstall(std::wstring packageId, std::wstring sourceId, std::unique_ptr<COMContext> context, bool isUpgrade)
     {
         std::unique_ptr<OrchestratorQueueItem> item = std::make_unique<OrchestratorQueueItem>(OrchestratorQueueItemId(std::move(packageId), std::move(sourceId)), std::move(context), isUpgrade ? PackageOperationType::Upgrade : PackageOperationType::Install);
         item->AddCommand(std::make_unique<::AppInstaller::CLI::COMDownloadCommand>(RootCommand::CommandName));
         item->AddCommand(std::make_unique<::AppInstaller::CLI::COMInstallCommand>(RootCommand::CommandName));
         return item;
     }
 
     std::unique_ptr<OrchestratorQueueItem> OrchestratorQueueItemFactory::CreateItemForUninstall(std::wstring packageId, std::wstring sourceId, std::unique_ptr<COMContext> context)
     {
         std::unique_ptr<OrchestratorQueueItem> item = std::make_unique<OrchestratorQueueItem>(OrchestratorQueueItemId(std::move(packageId), std::move(sourceId)), std::move(context), PackageOperationType::Uninstall);
         item->AddCommand(std::make_unique<::AppInstaller::CLI::COMUninstallCommand>(RootCommand::CommandName));
         return item;
     }
 
     std::unique_ptr<OrchestratorQueueItem> OrchestratorQueueItemFactory::CreateItemForSearch(std::wstring packageId, std::wstring sourceId, std::unique_ptr<COMContext> context)
     {
         std::unique_ptr<OrchestratorQueueItem> item = std::make_unique<OrchestratorQueueItem>(OrchestratorQueueItemId(std::move(packageId), std::move(sourceId)), std::move(context), PackageOperationType::Search);
         return item;
     }
 
     std::unique_ptr<OrchestratorQueueItem> OrchestratorQueueItemFactory::CreateItemForDownload(std::wstring packageId, std::wstring sourceId, std::unique_ptr<COMContext> context)
     {
         std::unique_ptr<OrchestratorQueueItem> item = std::make_unique<OrchestratorQueueItem>(OrchestratorQueueItemId(std::move(packageId), std::move(sourceId)), std::move(context), PackageOperationType::Download);
         item->AddCommand(std::make_unique<::AppInstaller::CLI::COMDownloadCommand>(RootCommand::CommandName));
         return item;
     }
 
     std::unique_ptr<OrchestratorQueueItem> OrchestratorQueueItemFactory::CreateItemForRepair(std::wstring packageId, std::wstring sourceId, std::unique_ptr<COMContext> context)
     {
         std::unique_ptr<OrchestratorQueueItem> item = std::make_unique<OrchestratorQueueItem>(OrchestratorQueueItemId(std::move(packageId), std::move(sourceId)), std::move(context), PackageOperationType::Repair);
         item->AddCommand(std::make_unique<::AppInstaller::CLI::COMRepairCommand>(RootCommand::CommandName));
         return item;
     }
 
     std::string_view ToString(OrchestratorQueueItemState state)
     {
         switch (state)
         {
         case OrchestratorQueueItemState::NotQueued: return "NotQueued";
         case OrchestratorQueueItemState::Queued: return "Queued";
         case OrchestratorQueueItemState::Running: return "Running";
         case OrchestratorQueueItemState::Cancelled: return "Cancelled";
         default: return "Unknown";
         }
     }
 }