Move timestamp pool reacquisition to an asynchronous worker queue, fix device_context race during destructor

1 files changed, 43 insertions, 28 deletions

diff --git a/src/timestamp_pool.hh b/src/timestamp_pool.hh
index d8ee359..bf3335d 100644
--- a/src/timestamp_pool.hh
+++ b/src/timestamp_pool.hh
@@ -2,20 +2,15 @@
 #define TIMESTAMP_POOL_HH_
 
 // The purpose of this file is to provide the definition of a 'timestamp pool'.
-// It manages blocks of timestamp query pools, hands them out when requested,
-// and allocates more when (if) we run out. It _should_ be thread safe.
-// Usage:
-//     1. Get handle with .acquire().
-//     2. Write start/end timestamp operations with the handle's pool and index
-//     into the provided command buffer. Will return nullopt if they're
-//     not yet available.
-//     3. Destruct the handle to return the key to the pool.
 
 #include <vulkan/utility/vk_dispatch_table.h>
 #include <vulkan/vulkan.hpp>
 
+#include <condition_variable>
+#include <deque>
 #include <memory>
 #include <mutex>
+#include <thread>
 #include <unordered_set>
 #include <vector>
 
@@ -26,19 +21,32 @@ namespace low_latency {
 class QueueContext;
 class DeviceContext;
 
+// A timestamp pool manages blocks of timestamp query pools, hands them out when
+// requested, and allocates more when (if) we run out. It _should_ be thread
+// safe.
+// Usage:
+//     1. Get handle with .acquire().
+//     2. Write start/end timestamp operations with the handle's pool and index
+//        into the provided command buffer.
+//     3. Grab the time, or wait until it's ready, using get_time or await_time
+//        respectively.
+//     4. Destruct the handle to return the key to the pool. The pool handles,
+//        via an async reaper thread, when the actual handle's contents can be
+//        reused as they must be alive until vulkan is done with them.
 class TimestampPool final {
   private:
     QueueContext& queue_context;
-    std::mutex mutex;
 
     // A chunk of data which is useful for making timestamp queries.
     // Allows association of an index to a query pool and command buffer.
     // We reuse these when they're released.
-    struct QueryChunk final {
+    class QueryChunk final {
+        friend class TimestampPool;
+
       private:
         static constexpr auto CHUNK_SIZE = 512u;
 
-      public:
+      private:
         struct QueryPoolOwner final {
           private:
             const QueueContext& queue_context;
@@ -55,10 +63,6 @@ class TimestampPool final {
           public:
             operator const VkQueryPool&() const { return this->query_pool; }
         };
-        std::unique_ptr<QueryPoolOwner> query_pool;
-
-        using free_indices_t = std::unordered_set<std::uint64_t>;
-        std::unique_ptr<free_indices_t> free_indices;
 
         struct CommandBuffersOwner final {
           private:
@@ -76,7 +80,11 @@ class TimestampPool final {
           public:
             VkCommandBuffer operator[](const std::size_t& i);
         };
+
+        std::unique_ptr<QueryPoolOwner> query_pool;
         std::unique_ptr<CommandBuffersOwner> command_buffers;
+        // A set of indices which are currently availabe in this chunk.
+        std::unordered_set<std::uint64_t> free_indices;
 
       public:
         QueryChunk(const QueueContext& queue_context);
@@ -86,19 +94,19 @@ class TimestampPool final {
         QueryChunk operator=(QueryChunk&&) = delete;
         ~QueryChunk();
     };
-    std::unordered_set<std::shared_ptr<QueryChunk>> query_chunks;
 
   public:
-    // A handle represents a VkCommandBuffer and a query index.
-    // Once the Handle goes out of scope, the query index will be returned
-    // to the parent pool.
+    // A handle represents a VkCommandBuffer and a query index. It can be used
+    // to attach timing information to submissions. Once the Handle destructs
+    // the query index will be returned to the parent pool - but crucially only
+    // when Vulkan is done with it.
     struct Handle final {
       private:
         friend class TimestampPool;
 
       private:
         TimestampPool& timestamp_pool;
-        const std::weak_ptr<QueryChunk> origin_chunk;
+        QueryChunk& query_chunk;
 
       public:
         const VkQueryPool query_pool;
@@ -106,13 +114,12 @@ class TimestampPool final {
         const VkCommandBuffer command_buffer;
 
       public:
-        Handle(TimestampPool& timestamp_pool,
-               const std::shared_ptr<QueryChunk>& origin_chunk,
+        Handle(TimestampPool& timestamp_pool, QueryChunk& query_chunk,
                const std::uint64_t& query_index);
         Handle(const Handle& handle) = delete;
-        Handle(Handle&&) = delete;
         Handle operator=(const Handle& handle) = delete;
-        Handle operator=(Handle&&) = delete;
+        Handle(Handle&&) = delete;
+        Handle& operator=(Handle&&) = delete;
         ~Handle();
 
       public:
@@ -125,11 +132,20 @@ class TimestampPool final {
 
         // Waits until the time is available and returns it.
         DeviceClock::time_point_t await_time();
-
-        // Calls get_time with the assumption it's already available.
-        DeviceClock::time_point_t get_time_required();
     };
 
+  private:
+    void do_reaper(const std::stop_token stoken);
+
+  private:
+    std::deque<Handle*> expiring_handles;
+    std::unordered_set<std::unique_ptr<QueryChunk>> query_chunks;
+
+    std::mutex mutex;
+    std::condition_variable_any cv;
+
+    std::jthread reaper_worker;
+
   public:
     TimestampPool(QueueContext& queue_context);
     TimestampPool(const TimestampPool&) = delete;
@@ -139,7 +155,6 @@ class TimestampPool final {
     ~TimestampPool();
 
   public:
-    // Hands out a Handle!
     std::shared_ptr<Handle> acquire();
 };