3 files changed, 42 insertions, 38 deletions

diff --git a/src/strategies/anti_lag/device_strategy.cc b/src/strategies/anti_lag/device_strategy.cc
index b391371..2d1d9ad 100644
--- a/src/strategies/anti_lag/device_strategy.cc
+++ b/src/strategies/anti_lag/device_strategy.cc
@@ -15,7 +15,7 @@ AntiLagDeviceStrategy::AntiLagDeviceStrategy(DeviceContext& device)
 AntiLagDeviceStrategy::~AntiLagDeviceStrategy() {}
 
 void AntiLagDeviceStrategy::notify_update(const VkAntiLagDataAMD& data) {
-    const auto lock = std::scoped_lock{this->mutex};
+    auto lock = std::unique_lock{this->mutex};
 
     this->is_enabled = !(data.mode == VK_ANTI_LAG_MODE_OFF_AMD);
 
@@ -37,11 +37,14 @@ void AntiLagDeviceStrategy::notify_update(const VkAntiLagDataAMD& data) {
         this->frame_index.reset();
         return;
     }
-
     // If we're at the input stage, start marking submissions as relevant.
     this->frame_index.emplace(data.pPresentationInfo->frameIndex);
 
-    { // Input stage needs to wait for all queue submissions to complete.
+    lock.unlock();
+    // We need to collect all queue submission and wait on them in this thread.
+    // Input stage needs to wait for all queue submissions to complete.
+    const auto queue_frame_spans = [&]() -> auto {
+        auto queue_frame_spans = std::vector<std::unique_ptr<FrameSpan>>{};
         const auto device_lock = std::shared_lock{this->device.mutex};
         for (const auto& iter : this->device.queues) {
             const auto& queue = iter.second;
@@ -50,15 +53,30 @@ void AntiLagDeviceStrategy::notify_update(const VkAntiLagDataAMD& data) {
                 dynamic_cast<AntiLagQueueStrategy*>(queue->strategy.get());
             assert(strategy);
 
-            strategy->await_complete();
+            // Grab it from the queue, don't hold the lock.
+            const auto queue_lock = std::scoped_lock{strategy->mutex};
+            queue_frame_spans.emplace_back(std::move(strategy->frame_span));
+            strategy->frame_span.reset();
+        }
+        return queue_frame_spans;
+    }();
+
+    // Wait on them and relock the mutex.
+    for (const auto& frame_span : queue_frame_spans) {
+        if (frame_span) { // Can still be null here.
+            frame_span->await_completed();
         }
     }
 
+    lock.lock();
+
     // We might need to wait a little more time to meet our frame limit.
     using namespace std::chrono;
     if (this->delay != 0us && this->previous_input_release.has_value()) {
+        lock.unlock();
         std::this_thread::sleep_until(*this->previous_input_release +
                                       this->delay);
+        lock.lock();
     }
 
     this->previous_input_release = steady_clock::now();
diff --git a/src/strategies/anti_lag/queue_strategy.cc b/src/strategies/anti_lag/queue_strategy.cc
index 27a9337..9c49e7c 100644
--- a/src/strategies/anti_lag/queue_strategy.cc
+++ b/src/strategies/anti_lag/queue_strategy.cc
@@ -12,7 +12,7 @@ AntiLagQueueStrategy::~AntiLagQueueStrategy() {}
 
 void AntiLagQueueStrategy::notify_submit(
     [[maybe_unused]] const VkSubmitInfo& submit,
-    std::unique_ptr<Submission> submission) {
+    std::shared_ptr<TimestampPool::Handle> handle) {
 
     const auto strategy =
         dynamic_cast<AntiLagDeviceStrategy*>(this->queue.device.strategy.get());
@@ -22,12 +22,16 @@ void AntiLagQueueStrategy::notify_submit(
     }
 
     const auto lock = std::scoped_lock(this->mutex);
-    this->pending_submissions.push_back(std::move(submission));
+    if (this->frame_span) {
+        this->frame_span->update(std::move(handle));
+    } else {
+        this->frame_span = std::make_unique<FrameSpan>(std::move(handle));
+    }
 }
 
 void AntiLagQueueStrategy::notify_submit(
     [[maybe_unused]] const VkSubmitInfo2& submit,
-    std::unique_ptr<Submission> submission) {
+    std::shared_ptr<TimestampPool::Handle> handle) {
 
     const auto strategy =
         dynamic_cast<AntiLagDeviceStrategy*>(this->queue.device.strategy.get());
@@ -37,26 +41,11 @@ void AntiLagQueueStrategy::notify_submit(
     }
 
     const auto lock = std::scoped_lock(this->mutex);
-    this->pending_submissions.push_back(std::move(submission));
-}
-
-void AntiLagQueueStrategy::await_complete() {
-
-    // Grab submissions while under a lock.
-    const auto submissions = [&]() -> std::deque<std::unique_ptr<Submission>> {
-        const auto lock = std::scoped_lock{this->mutex};
-
-        auto submissions = std::move(this->pending_submissions);
-        this->pending_submissions.clear();
-        return submissions;
-    }();
-
-    // Wait for completion on the last submission.
-    if (submissions.empty()) {
-        return;
+    if (this->frame_span) {
+        this->frame_span->update(std::move(handle));
+    } else {
+        this->frame_span = std::make_unique<FrameSpan>(std::move(handle));
     }
-    const auto& last = submissions.back();
-    last->handle->await_end_time();
 }
 
 // Stub - AntiLag doesn't care about presents.
diff --git a/src/strategies/anti_lag/queue_strategy.hh b/src/strategies/anti_lag/queue_strategy.hh
index 37c44a5..b1ae3e6 100644
--- a/src/strategies/anti_lag/queue_strategy.hh
+++ b/src/strategies/anti_lag/queue_strategy.hh
@@ -3,7 +3,7 @@
 
 #include "strategies/queue_strategy.hh"
 
-#include <deque>
+#include "frame_span.hh"
 #include <memory>
 #include <mutex>
 
@@ -12,25 +12,22 @@ namespace low_latency {
 class QueueContext;
 
 class AntiLagQueueStrategy final : public QueueStrategy {
-  private:
+  public:
     std::mutex mutex;
-    std::deque<std::unique_ptr<Submission>> pending_submissions;
+    std::unique_ptr<FrameSpan> frame_span; // Null represents no work.
 
   public:
     AntiLagQueueStrategy(QueueContext& queue);
     virtual ~AntiLagQueueStrategy();
 
   public:
-    virtual void notify_submit(const VkSubmitInfo& submit,
-                               std::unique_ptr<Submission> submission) override;
-    virtual void notify_submit(const VkSubmitInfo2& submit,
-                               std::unique_ptr<Submission> submission) override;
+    virtual void
+    notify_submit(const VkSubmitInfo& submit,
+                  std::shared_ptr<TimestampPool::Handle> handle) override;
+    virtual void
+    notify_submit(const VkSubmitInfo2& submit,
+                  std::shared_ptr<TimestampPool::Handle> handle) override;
     virtual void notify_present(const VkPresentInfoKHR& present) override;
-
-  public:
-    // Wait for all pending submissions to complete. Resets pending submissions
-    // once done.
-    void await_complete();
 };
 
 } // namespace low_latency