From 3adf9e596745abe6121e83f93c120bfbc56bc9b6 Mon Sep 17 00:00:00 2001
From: Cameron Gutman <aicommander@gmail.com>
Date: Wed, 3 Jan 2024 16:08:46 -0600
Subject: [PATCH] Don't busy loop during wlgrab capture

---
 src/platform/linux/wayland.cpp | 33 +++++++++++++++++++++++++++++++++
 src/platform/linux/wayland.h   |  4 ++++
 src/platform/linux/wlgrab.cpp  |  6 +++---
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/src/platform/linux/wayland.cpp b/src/platform/linux/wayland.cpp
index bc1e416f..b762492b 100644
--- a/src/platform/linux/wayland.cpp
+++ b/src/platform/linux/wayland.cpp
@@ -2,6 +2,7 @@
  * @file src/platform/linux/wayland.cpp
  * @brief todo
  */
+#include <poll.h>
 #include <wayland-client.h>
 #include <wayland-util.h>
 
@@ -61,6 +62,38 @@ namespace wl {
     wl_display_roundtrip(display_internal.get());
   }
 
+  /**
+   * @brief Waits up to the specified timeout to dispatch new events on the wl_display.
+   * @param timeout The timeout in milliseconds.
+   * @return true if new events were dispatched or false if the timeout expired.
+   */
+  bool
+  display_t::dispatch(std::chrono::milliseconds timeout) {
+    // Check if any events are queued already. If not, flush
+    // outgoing events, and prepare to wait for readability.
+    if (wl_display_prepare_read(display_internal.get()) == 0) {
+      wl_display_flush(display_internal.get());
+
+      // Wait for an event to come in
+      struct pollfd pfd = {};
+      pfd.fd = wl_display_get_fd(display_internal.get());
+      pfd.events = POLLIN;
+      if (poll(&pfd, 1, timeout.count()) == 1 && (pfd.revents & POLLIN)) {
+        // Read the new event(s)
+        wl_display_read_events(display_internal.get());
+      }
+      else {
+        // We timed out, so unlock the queue now
+        wl_display_cancel_read(display_internal.get());
+        return false;
+      }
+    }
+
+    // Dispatch any existing or new pending events
+    wl_display_dispatch_pending(display_internal.get());
+    return true;
+  }
+
   wl_registry *
   display_t::registry() {
     return wl_display_get_registry(display_internal.get());
diff --git a/src/platform/linux/wayland.h b/src/platform/linux/wayland.h
index e4c1c1a9..062aa65f 100644
--- a/src/platform/linux/wayland.h
+++ b/src/platform/linux/wayland.h
@@ -206,6 +206,10 @@ namespace wl {
     void
     roundtrip();
 
+    // Wait up to the timeout to read and dispatch new events
+    bool
+    dispatch(std::chrono::milliseconds timeout);
+
     // Get the registry associated with the display
     // No need to manually free the registry
     wl_registry *
diff --git a/src/platform/linux/wlgrab.cpp b/src/platform/linux/wlgrab.cpp
index 8f3ef299..5ba0d9e2 100644
--- a/src/platform/linux/wlgrab.cpp
+++ b/src/platform/linux/wlgrab.cpp
@@ -87,11 +87,11 @@ namespace wl {
     snapshot(const pull_free_image_cb_t &pull_free_image_cb, std::shared_ptr<platf::img_t> &img_out, std::chrono::milliseconds timeout, bool cursor) {
       auto to = std::chrono::steady_clock::now() + timeout;
 
+      // Dispatch events until we get a new frame or the timeout expires
       dmabuf.listen(interface.dmabuf_manager, output, cursor);
       do {
-        display.roundtrip();
-
-        if (to < std::chrono::steady_clock::now()) {
+        auto remaining_time_ms = std::chrono::duration_cast<std::chrono::milliseconds>(to - std::chrono::steady_clock::now());
+        if (remaining_time_ms.count() < 0 || !display.dispatch(remaining_time_ms)) {
           return platf::capture_e::timeout;
         }
       } while (dmabuf.status == dmabuf_t::WAITING);