[Impeller] vk::CommandPool resets via ResourceManager

impeller/renderer/backend/vulkan/blit_command_vk_unittests.cc

@@ -2,7 +2,7 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-#include "flutter/testing/testing.h"
+#include "flutter/testing/testing.h"  // IWYU pragma: keep
 #include "impeller/renderer/backend/vulkan/blit_command_vk.h"
 #include "impeller/renderer/backend/vulkan/command_encoder_vk.h"
 #include "impeller/renderer/backend/vulkan/test/mock_vulkan.h"
@@ -12,7 +12,6 @@ namespace testing {
 
 TEST(BlitCommandVkTest, BlitCopyTextureToTextureCommandVK) {
   auto context = CreateMockVulkanContext();
-  auto pool = CommandPoolVK::GetThreadLocal(context.get());
   auto encoder = std::make_unique<CommandEncoderFactoryVK>(context)->Create();
   BlitCopyTextureToTextureCommandVK cmd;
   cmd.source = context->GetResourceAllocator()->CreateTexture({

impeller/renderer/backend/vulkan/command_encoder_vk.cc

@@ -5,7 +5,6 @@
 #include "impeller/renderer/backend/vulkan/command_encoder_vk.h"
 
 #include "flutter/fml/closure.h"
-#include "flutter/fml/trace_event.h"
 #include "impeller/renderer/backend/vulkan/context_vk.h"
 #include "impeller/renderer/backend/vulkan/fence_waiter_vk.h"
 #include "impeller/renderer/backend/vulkan/texture_vk.h"
@@ -21,7 +20,7 @@ class TrackedObjectsVK {
     if (!pool) {
       return;
     }
-    auto buffer = pool->CreateGraphicsCommandBuffer();
+    auto buffer = pool->CreateBuffer();
     if (!buffer) {
       return;
     }
@@ -34,7 +33,7 @@ class TrackedObjectsVK {
     if (!buffer_) {
       return;
     }
-    pool_->CollectGraphicsCommandBuffer(std::move(buffer_));
+    pool_->CollectBuffer(std::move(buffer_));
   }
 
   bool IsValid() const { return is_valid_; }
@@ -105,7 +104,11 @@ std::shared_ptr<CommandEncoderVK> CommandEncoderFactoryVK::Create() {
     return nullptr;
   }
   auto& context_vk = ContextVK::Cast(*context);
-  auto tls_pool = CommandPoolVK::GetThreadLocal(&context_vk);
+  auto recycler = context_vk.GetCommandPoolRecycler();
+  if (!recycler) {
+    return nullptr;
+  }
+  auto tls_pool = recycler->Get();
   if (!tls_pool) {
     return nullptr;
   }

impeller/renderer/backend/vulkan/command_pool_vk.cc

@@ -3,166 +3,190 @@
 // found in the LICENSE file.
 
 #include "impeller/renderer/backend/vulkan/command_pool_vk.h"
+#include <memory>
+#include <optional>
+#include <utility>
+#include "fml/macros.h"
+#include "fml/thread_local.h"
+#include "fml/trace_event.h"
+#include "impeller/renderer/backend/vulkan/resource_manager_vk.h"
+#include "vulkan/vulkan_structs.hpp"
 
-#include <map>
-#include <unordered_map>
-#include <vector>
+namespace impeller {
 
-#include "flutter/fml/thread_local.h"
-#include "impeller/base/thread.h"
-#include "impeller/renderer/backend/vulkan/context_vk.h"
+// Holds the command pool in a background thread, recyling it when not in use.
+class BackgroundCommandPoolVK {
+ public:
+  BackgroundCommandPoolVK(BackgroundCommandPoolVK&&) = default;
+
+  explicit BackgroundCommandPoolVK(
+      vk::UniqueCommandPool&& pool,
+      std::vector<vk::UniqueCommandBuffer>&& buffers,
+      std::weak_ptr<CommandPoolRecyclerVK> recycler)
+      : pool_(std::move(pool)),
+        buffers_(std::move(buffers)),
+        recycler_(std::move(recycler)) {}
+
+  ~BackgroundCommandPoolVK() {
+    auto const recycler = recycler_.lock();
+    if (!recycler) {
+      return;
+    }
+    recycler->Reclaim(std::move(pool_));
+  }
 
-namespace impeller {
+ private:
+  FML_DISALLOW_COPY_AND_ASSIGN(BackgroundCommandPoolVK);
 
-using CommandPoolMap = std::map<uint64_t, std::shared_ptr<CommandPoolVK>>;
-FML_THREAD_LOCAL fml::ThreadLocalUniquePtr<CommandPoolMap> tls_command_pool;
+  vk::UniqueCommandPool pool_;
 
-static Mutex g_all_pools_mutex;
-static std::unordered_map<const ContextVK*,
-                          std::vector<std::weak_ptr<CommandPoolVK>>>
-    g_all_pools IPLR_GUARDED_BY(g_all_pools_mutex);
+  // These are retained otherwise it is a Vulkan thread-safety violation
+  // because the buffers and the originating pool do not belong to the same
+  // thread.
+  std::vector<vk::UniqueCommandBuffer> buffers_;
+  std::weak_ptr<CommandPoolRecyclerVK> recycler_;
+};
 
-std::shared_ptr<CommandPoolVK> CommandPoolVK::GetThreadLocal(
-    const ContextVK* context) {
+CommandPoolVK::~CommandPoolVK() {
+  auto const context = context_.lock();
   if (!context) {
-    return nullptr;
-  }
-  if (tls_command_pool.get() == nullptr) {
-    tls_command_pool.reset(new CommandPoolMap());
-  }
-  CommandPoolMap& pool_map = *tls_command_pool.get();
-  auto found = pool_map.find(context->GetHash());
-  if (found != pool_map.end() && found->second->IsValid()) {
-    return found->second;
-  }
-  auto pool = std::shared_ptr<CommandPoolVK>(new CommandPoolVK(context));
-  if (!pool->IsValid()) {
-    return nullptr;
+    return;
   }
-  pool_map[context->GetHash()] = pool;
-  {
-    Lock pool_lock(g_all_pools_mutex);
-    g_all_pools[context].push_back(pool);
+  auto const recycler = context->GetCommandPoolRecycler();
+  if (!recycler) {
+    return;
   }
-  return pool;
+  UniqueResourceVKT<BackgroundCommandPoolVK> pool(
+      context->GetResourceManager(),
+      BackgroundCommandPoolVK(std::move(pool_), std::move(collected_buffers_),
+                              recycler));
 }
 
-void CommandPoolVK::ClearAllPools(const ContextVK* context) {
-  if (tls_command_pool.get()) {
-    tls_command_pool.get()->erase(context->GetHash());
-  }
-  Lock pool_lock(g_all_pools_mutex);
-  if (auto found = g_all_pools.find(context); found != g_all_pools.end()) {
-    for (auto& weak_pool : found->second) {
-      auto pool = weak_pool.lock();
-      if (!pool) {
-        // The pool has already died because the thread died.
-        continue;
-      }
-      // The pool is reset but its reference in the TLS map remains till the
-      // thread dies.
-      pool->Reset();
-    }
-    g_all_pools.erase(found);
+// TODO(matanlurey): Return a status_or<> instead of {} when we have one.
+vk::UniqueCommandBuffer CommandPoolVK::CreateBuffer() {
+  auto const context = context_.lock();
+  if (!context) {
+    return {};
   }
-}
 
-CommandPoolVK::CommandPoolVK(const ContextVK* context)
-    : owner_id_(std::this_thread::get_id()) {
-  vk::CommandPoolCreateInfo pool_info;
-
-  pool_info.queueFamilyIndex = context->GetGraphicsQueue()->GetIndex().family;
-  pool_info.flags = vk::CommandPoolCreateFlagBits::eTransient |
-                    vk::CommandPoolCreateFlagBits::eResetCommandBuffer;
-  auto pool = context->GetDevice().createCommandPoolUnique(pool_info);
-  if (pool.result != vk::Result::eSuccess) {
-    return;
+  auto const device = context->GetDevice();
+  vk::CommandBufferAllocateInfo info;
+  info.setCommandPool(pool_.get());
+  info.setCommandBufferCount(1u);
+  info.setLevel(vk::CommandBufferLevel::ePrimary);
+  auto [result, buffers] = device.allocateCommandBuffersUnique(info);
+  if (result != vk::Result::eSuccess) {
+    return {};
   }
-
-  device_holder_ = context->GetDeviceHolder();
-  graphics_pool_ = std::move(pool.value);
-  is_valid_ = true;
+  return std::move(buffers[0]);
 }
 
-CommandPoolVK::~CommandPoolVK() = default;
-
-bool CommandPoolVK::IsValid() const {
-  return is_valid_;
+void CommandPoolVK::CollectBuffer(vk::UniqueCommandBuffer&& buffer) {
+  collected_buffers_.push_back(std::move(buffer));
 }
 
-void CommandPoolVK::Reset() {
-  {
-    Lock lock(buffers_to_collect_mutex_);
-    graphics_pool_.reset();
+// Associates a resource with a thread and context.
+using CommandPoolMap =
+    std::unordered_map<uint64_t, std::shared_ptr<CommandPoolVK>>;
+FML_THREAD_LOCAL fml::ThreadLocalUniquePtr<CommandPoolMap> resources_;
 
-    // When the command pool is destroyed, all of its command buffers are freed.
-    // Handles allocated from that pool are now invalid and must be discarded.
-    for (vk::UniqueCommandBuffer& buffer : buffers_to_collect_) {
-      buffer.release();
-    }
-    buffers_to_collect_.clear();
+// 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();
+  if (!strong_context) {
+    return nullptr;
   }
 
-  for (vk::UniqueCommandBuffer& buffer : recycled_buffers_) {
-    buffer.release();
+  // If there is a resource in used for this thread and context, return it.
+  auto resources = resources_.get();
+  if (!resources) {
+    resources = new CommandPoolMap();
+    resources_.reset(resources);
+  }
+  auto map = *resources;
+  auto const hash = strong_context->GetHash();
+  auto const it = map.find(hash);
+  if (it != map.end()) {
+    return it->second;
   }
-  recycled_buffers_.clear();
 
-  is_valid_ = false;
+  // Otherwise, create a new resource and return it.
+  auto pool = Create();
+  if (!pool) {
+    return nullptr;
+  }
+
+  auto const resource =
+      std::make_shared<CommandPoolVK>(std::move(*pool), context_);
+  map[hash] = resource;
+  return resource;
 }
 
-vk::UniqueCommandBuffer CommandPoolVK::CreateGraphicsCommandBuffer() {
-  std::shared_ptr<const DeviceHolder> strong_device = device_holder_.lock();
-  if (!strong_device) {
-    return {};
-  }
-  FML_DCHECK(std::this_thread::get_id() == owner_id_);
-  {
-    Lock lock(buffers_to_collect_mutex_);
-    GarbageCollectBuffersIfAble();
+// TODO(matanlurey): Return a status_or<> instead of nullopt when we have one.
+std::optional<vk::UniqueCommandPool> CommandPoolRecyclerVK::Create() {
+  // If we can reuse a command pool, do so.
+  if (auto pool = Reuse()) {
+    return pool;
   }
 
-  if (!recycled_buffers_.empty()) {
-    vk::UniqueCommandBuffer result = std::move(recycled_buffers_.back());
-    recycled_buffers_.pop_back();
-    return result;
+  // Otherwise, create a new one.
+  auto context = context_.lock();
+  if (!context) {
+    return std::nullopt;
   }
+  vk::CommandPoolCreateInfo info;
+  info.setQueueFamilyIndex(context->GetGraphicsQueue()->GetIndex().family);
+  info.setFlags(vk::CommandPoolCreateFlagBits::eTransient);
 
-  vk::CommandBufferAllocateInfo alloc_info;
-  alloc_info.commandPool = graphics_pool_.get();
-  alloc_info.commandBufferCount = 1u;
-  alloc_info.level = vk::CommandBufferLevel::ePrimary;
-  auto [result, buffers] =
-      strong_device->GetDevice().allocateCommandBuffersUnique(alloc_info);
+  auto device = context->GetDevice();
+  auto [result, pool] = device.createCommandPoolUnique(info);
   if (result != vk::Result::eSuccess) {
-    return {};
+    return std::nullopt;
   }
-  return std::move(buffers[0]);
+  return std::move(pool);
 }
 
-void CommandPoolVK::CollectGraphicsCommandBuffer(
-    vk::UniqueCommandBuffer buffer) {
-  Lock lock(buffers_to_collect_mutex_);
-  if (!graphics_pool_) {
-    // If the command pool has already been destroyed, then its command buffers
-    // have been freed and are now invalid.
-    buffer.release();
+std::optional<vk::UniqueCommandPool> CommandPoolRecyclerVK::Reuse() {
+  // If there are no recycled pools, return nullopt.
+  Lock _(recycled_mutex_);
+  if (recycled_.empty()) {
+    return std::nullopt;
   }
-  buffers_to_collect_.emplace_back(std::move(buffer));
-  GarbageCollectBuffersIfAble();
+
+  // Otherwise, remove and return a recycled pool.
+  auto pool = std::move(recycled_.back());
+  recycled_.pop_back();
+  return std::move(pool);
 }
 
-void CommandPoolVK::GarbageCollectBuffersIfAble() {
-  if (std::this_thread::get_id() != owner_id_) {
+void CommandPoolRecyclerVK::Reclaim(vk::UniqueCommandPool&& pool) {
+  TRACE_EVENT0("impeller", "ReclaimCommandPool");
+  // FIXME: Assert that this is called on a background thread.
+
+  // Reset the pool on a background thread.
+  auto strong_context = context_.lock();
+  if (!strong_context) {
     return;
   }
+  auto device = strong_context->GetDevice();
+  device.resetCommandPool(pool.get());
 
-  for (auto& buffer : buffers_to_collect_) {
-    buffer->reset();
-    recycled_buffers_.emplace_back(std::move(buffer));
-  }
+  // Move the pool to the recycled list.
+  Lock _(recycled_mutex_);
+  recycled_.push_back(std::move(pool));
+}
+
+CommandPoolRecyclerVK::~CommandPoolRecyclerVK() {
+  // Ensure all recycled pools are reclaimed before this is destroyed.
+  Recycle();
+}
 
-  buffers_to_collect_.clear();
+void CommandPoolRecyclerVK::Recycle() {
+  auto const resources = resources_.get();
+  if (!resources) {
+    return;
+  }
+  resources->clear();
 }
 
 }  // namespace impeller