[Impeller] Maintain a global map of each context's currently active thread-local command pools (#169548)

The Impeller Vulkan back end creates a thread-local map of contexts to
CommandPoolVK instances for each thread that uses Vulkan. This allows a
thread to obtain the CommandPoolVK that is currently in use for a given
context.

When a context is shut down, the Vulkan resources used by each thread's
local CommandPoolVK for that context must be freed. To do this,
CommandPoolVK maintains a global map of CommandPoolVK instances.

Prior to this PR Impeller was appending to the context's pool list in
the global map each time a new CommandPoolVK was created. In the
original implementation this worked because Impeller was only creating
one CommandPoolVK per thread for a given context.

However, CommandPoolVK later adopted a recycling scheme where each frame
creates a new CommandPoolVK instance that acquires a Vulkan command pool
from the CommandPoolRecyclerVK. So inserting every CommandPoolVK into
the global map will cause the global map to grow unbounded.

This PR changes the structure of the global map. The global map will now
associate each context with a map of thread IDs to the CommandPoolVK
that is currently placed in the thread's local storage. When a thread
calls CommandPoolRecyclerVK::Dispose to clear its thread-local
CommandPoolVK for a context, the corresponding entry is also removed
from the global map.

Fixes https://github.com/flutter/flutter/issues/169208

engine/src/flutter/impeller/golden_tests/golden_playground_test_mac.cc

@@ -132,6 +132,11 @@ void GoldenPlaygroundTest::SetTypographerContext(
 
 void GoldenPlaygroundTest::TearDown() {
   ASSERT_FALSE(dlopen("/usr/local/lib/libMoltenVK.dylib", RTLD_NOLOAD));
+
+  auto context = GetContext();
+  if (context) {
+    context->DisposeThreadLocalCachedResources();
+  }
 }
 
 namespace {
@@ -280,6 +285,9 @@ RuntimeStage::Map GoldenPlaygroundTest::OpenAssetAsRuntimeStage(
 }
 
 std::shared_ptr<Context> GoldenPlaygroundTest::GetContext() const {
+  if (!pimpl_->screenshotter) {
+    return nullptr;
+  }
   return pimpl_->screenshotter->GetPlayground().GetContext();
 }
 

engine/src/flutter/impeller/renderer/backend/vulkan/command_pool_vk.cc

@@ -171,10 +171,23 @@ static thread_local std::unique_ptr<CommandPoolMap> tls_command_pool_map;
 // with that context.
 static Mutex g_all_pools_map_mutex;
 static std::unordered_map<
-    const ContextVK*,
-    std::vector<std::weak_ptr<CommandPoolVK>>> g_all_pools_map
+    uint64_t,
+    std::unordered_map<std::thread::id,
+                       std::weak_ptr<CommandPoolVK>>> g_all_pools_map
     IPLR_GUARDED_BY(g_all_pools_map_mutex);
 
+CommandPoolRecyclerVK::CommandPoolRecyclerVK(
+    const std::shared_ptr<ContextVK>& context)
+    : context_(context), context_hash_(context->GetHash()) {}
+
+// Visible for testing.
+// Returns the number of pools in g_all_pools_map for the given context.
+int CommandPoolRecyclerVK::GetGlobalPoolCount(const ContextVK& context) {
+  Lock all_pools_lock(g_all_pools_map_mutex);
+  auto it = g_all_pools_map.find(context.GetHash());
+  return it != g_all_pools_map.end() ? it->second.size() : 0;
+}
+
 // TODO(matanlurey): Return a status_or<> instead of nullptr when we have one.
 std::shared_ptr<CommandPoolVK> CommandPoolRecyclerVK::Get() {
   auto const strong_context = context_.lock();
@@ -187,8 +200,7 @@ std::shared_ptr<CommandPoolVK> CommandPoolRecyclerVK::Get() {
     tls_command_pool_map.reset(new CommandPoolMap());
   }
   CommandPoolMap& pool_map = *tls_command_pool_map.get();
-  auto const hash = strong_context->GetHash();
-  auto const it = pool_map.find(hash);
+  auto const it = pool_map.find(context_hash_);
   if (it != pool_map.end()) {
     return it->second;
   }
@@ -201,11 +213,11 @@ std::shared_ptr<CommandPoolVK> CommandPoolRecyclerVK::Get() {
 
   auto const resource = std::make_shared<CommandPoolVK>(
       std::move(data->pool), std::move(data->buffers), context_);
-  pool_map.emplace(hash, resource);
+  pool_map.emplace(context_hash_, resource);
 
   {
     Lock all_pools_lock(g_all_pools_map_mutex);
-    g_all_pools_map[strong_context.get()].push_back(resource);
+    g_all_pools_map[context_hash_][std::this_thread::get_id()] = resource;
   }
 
   return resource;
@@ -275,30 +287,33 @@ void CommandPoolRecyclerVK::Reclaim(
       RecycledData{.pool = std::move(pool), .buffers = std::move(buffers)});
 }
 
-CommandPoolRecyclerVK::~CommandPoolRecyclerVK() {
-  // Ensure all recycled pools are reclaimed before this is destroyed.
-  Dispose();
-}
-
 void CommandPoolRecyclerVK::Dispose() {
   CommandPoolMap* pool_map = tls_command_pool_map.get();
   if (pool_map) {
-    pool_map->clear();
+    pool_map->erase(context_hash_);
+  }
+
+  {
+    Lock all_pools_lock(g_all_pools_map_mutex);
+    auto found = g_all_pools_map.find(context_hash_);
+    if (found != g_all_pools_map.end()) {
+      found->second.erase(std::this_thread::get_id());
+    }
   }
 }
 
-void CommandPoolRecyclerVK::DestroyThreadLocalPools(const ContextVK* context) {
+void CommandPoolRecyclerVK::DestroyThreadLocalPools() {
   // Delete the context's entry in this thread's command pool map.
   if (tls_command_pool_map.get()) {
-    tls_command_pool_map.get()->erase(context->GetHash());
+    tls_command_pool_map.get()->erase(context_hash_);
   }
 
   // Destroy all other thread-local CommandPoolVK instances associated with
   // this context.
   Lock all_pools_lock(g_all_pools_map_mutex);
-  auto found = g_all_pools_map.find(context);
+  auto found = g_all_pools_map.find(context_hash_);
   if (found != g_all_pools_map.end()) {
-    for (auto& weak_pool : found->second) {
+    for (auto& [thread_id, weak_pool] : found->second) {
       auto pool = weak_pool.lock();
       if (!pool) {
         continue;

engine/src/flutter/impeller/renderer/backend/vulkan/command_pool_vk.h

@@ -103,8 +103,6 @@ class CommandPoolVK final {
 class CommandPoolRecyclerVK final
     : public std::enable_shared_from_this<CommandPoolRecyclerVK> {
  public:
-  ~CommandPoolRecyclerVK();
-
   /// A unique command pool and zero or more recycled command buffers.
   struct RecycledData {
     vk::UniqueCommandPool pool;
@@ -112,16 +110,13 @@ class CommandPoolRecyclerVK final
   };
 
   /// @brief      Clean up resources held by all per-thread command pools
-  ///             associated with the given context.
-  ///
-  /// @param[in]  context The context.
-  static void DestroyThreadLocalPools(const ContextVK* context);
+  ///             associated with the context.
+  void DestroyThreadLocalPools();
 
   /// @brief      Creates a recycler for the given |ContextVK|.
   ///
   /// @param[in]  context The context to create the recycler for.
-  explicit CommandPoolRecyclerVK(std::weak_ptr<ContextVK> context)
-      : context_(std::move(context)) {}
+  explicit CommandPoolRecyclerVK(const std::shared_ptr<ContextVK>& context);
 
   /// @brief      Gets a command pool for the current thread.
   ///
@@ -137,11 +132,15 @@ class CommandPoolRecyclerVK final
                std::vector<vk::UniqueCommandBuffer>&& buffers,
                bool should_trim = false);
 
-  /// @brief      Clears all recycled command pools to let them be reclaimed.
+  /// @brief      Clears this context's thread-local command pool.
   void Dispose();
 
+  // Visible for testing.
+  static int GetGlobalPoolCount(const ContextVK& context);
+
  private:
   std::weak_ptr<ContextVK> context_;
+  uint64_t context_hash_;
 
   Mutex recycled_mutex_;
   std::vector<RecycledData> recycled_ IPLR_GUARDED_BY(recycled_mutex_);

engine/src/flutter/impeller/renderer/backend/vulkan/command_pool_vk_unittests.cc

@@ -228,5 +228,24 @@ TEST(CommandPoolRecyclerVKTest, ExtraCommandBufferAllocationsTriggerTrim) {
   context->Shutdown();
 }
 
+TEST(CommandPoolRecyclerVKTest, RecyclerGlobalPoolMapSize) {
+  auto context = MockVulkanContextBuilder().Build();
+  auto const recycler = context->GetCommandPoolRecycler();
+
+  // The global pool list for this context should initially be empty.
+  EXPECT_EQ(CommandPoolRecyclerVK::GetGlobalPoolCount(*context), 0);
+
+  // Creating a pool for this thread should insert the pool into the global map.
+  auto pool = recycler->Get();
+  EXPECT_EQ(CommandPoolRecyclerVK::GetGlobalPoolCount(*context), 1);
+
+  // Disposing this thread's pool should remove it from the global map.
+  pool.reset();
+  recycler->Dispose();
+  EXPECT_EQ(CommandPoolRecyclerVK::GetGlobalPoolCount(*context), 0);
+
+  context->Shutdown();
+}
+
 }  // namespace testing
 }  // namespace impeller

engine/src/flutter/impeller/renderer/backend/vulkan/context_vk.cc

@@ -134,7 +134,9 @@ ContextVK::~ContextVK() {
   if (device_holder_ && device_holder_->device) {
     [[maybe_unused]] auto result = device_holder_->device->waitIdle();
   }
-  CommandPoolRecyclerVK::DestroyThreadLocalPools(this);
+  if (command_pool_recycler_) {
+    command_pool_recycler_->DestroyThreadLocalPools();
+  }
 }
 
 Context::BackendType ContextVK::GetBackendType() const {
@@ -421,7 +423,7 @@ void ContextVK::Setup(Settings settings) {
   }
 
   auto command_pool_recycler =
-      std::make_shared<CommandPoolRecyclerVK>(weak_from_this());
+      std::make_shared<CommandPoolRecyclerVK>(shared_from_this());
   if (!command_pool_recycler) {
     VALIDATION_LOG << "Could not create command pool recycler.";
     return;