From cbbe49cdf014e268a609e55e19f867329217db80 Mon Sep 17 00:00:00 2001
From: Aliaksandr Kalenik <kalenik.aliaksandr@gmail.com>
Date: Fri, 21 Jun 2024 19:15:32 +0300
Subject: [PATCH] WebContent: Move screenshot requests handling into repaint
 loop

This change ensures that if a screenshot is requested between teardown
of paintable tree and relayout, we will wait for layout to complete
before taking a screenshot.
---
 .../WebContent/ConnectionFromClient.cpp       | 27 ++-----------------
 Userland/Services/WebContent/PageClient.cpp   | 26 ++++++++++++++++++
 Userland/Services/WebContent/PageClient.h     |  7 +++++
 3 files changed, 35 insertions(+), 25 deletions(-)

diff --git a/Userland/Services/WebContent/ConnectionFromClient.cpp b/Userland/Services/WebContent/ConnectionFromClient.cpp
index 2bdf6cae636..f4d1671129b 100644
--- a/Userland/Services/WebContent/ConnectionFromClient.cpp
+++ b/Userland/Services/WebContent/ConnectionFromClient.cpp
@@ -754,19 +754,7 @@ void ConnectionFromClient::take_document_screenshot(u64 page_id)
     if (!page.has_value())
         return;
 
-    auto* document = page->page().top_level_browsing_context().active_document();
-    if (!document || !document->document_element()) {
-        async_did_take_screenshot(page_id, {});
-        return;
-    }
-
-    auto const& content_size = page->content_size();
-    Web::DevicePixelRect rect { { 0, 0 }, content_size };
-
-    auto bitmap = Gfx::Bitmap::create(Gfx::BitmapFormat::BGRA8888, rect.size().to_type<int>()).release_value_but_fixme_should_propagate_errors();
-    page->paint(rect, *bitmap);
-
-    async_did_take_screenshot(page_id, bitmap->to_shareable_bitmap());
+    page->queue_screenshot_task({});
 }
 
 void ConnectionFromClient::take_dom_node_screenshot(u64 page_id, i32 node_id)
@@ -775,18 +763,7 @@ void ConnectionFromClient::take_dom_node_screenshot(u64 page_id, i32 node_id)
     if (!page.has_value())
         return;
 
-    auto* dom_node = Web::DOM::Node::from_unique_id(node_id);
-    if (!dom_node || !dom_node->paintable_box()) {
-        async_did_take_screenshot(page_id, {});
-        return;
-    }
-
-    auto rect = page->page().enclosing_device_rect(dom_node->paintable_box()->absolute_border_box_rect());
-
-    auto bitmap = Gfx::Bitmap::create(Gfx::BitmapFormat::BGRA8888, rect.size().to_type<int>()).release_value_but_fixme_should_propagate_errors();
-    page->paint(rect, *bitmap, { .paint_overlay = Web::PaintOptions::PaintOverlay::No });
-
-    async_did_take_screenshot(page_id, bitmap->to_shareable_bitmap());
+    page->queue_screenshot_task(node_id);
 }
 
 Messages::WebContentServer::DumpGcGraphResponse ConnectionFromClient::dump_gc_graph(u64)
diff --git a/Userland/Services/WebContent/PageClient.cpp b/Userland/Services/WebContent/PageClient.cpp
index 476d6192d56..a685f44c59a 100644
--- a/Userland/Services/WebContent/PageClient.cpp
+++ b/Userland/Services/WebContent/PageClient.cpp
@@ -196,6 +196,26 @@ Web::Layout::Viewport* PageClient::layout_root()
 
 void PageClient::paint_next_frame()
 {
+    while (!m_screenshot_tasks.is_empty()) {
+        auto task = m_screenshot_tasks.dequeue();
+        if (task.node_id.has_value()) {
+            auto* dom_node = Web::DOM::Node::from_unique_id(*task.node_id);
+            if (!dom_node || !dom_node->paintable_box()) {
+                client().async_did_take_screenshot(m_id, {});
+                continue;
+            }
+            auto rect = page().enclosing_device_rect(dom_node->paintable_box()->absolute_border_box_rect());
+            auto bitmap = Gfx::Bitmap::create(Gfx::BitmapFormat::BGRA8888, rect.size().to_type<int>()).release_value_but_fixme_should_propagate_errors();
+            paint(rect, *bitmap, { .paint_overlay = Web::PaintOptions::PaintOverlay::No });
+            client().async_did_take_screenshot(m_id, bitmap->to_shareable_bitmap());
+        } else {
+            Web::DevicePixelRect rect { { 0, 0 }, content_size() };
+            auto bitmap = Gfx::Bitmap::create(Gfx::BitmapFormat::BGRA8888, rect.size().to_type<int>()).release_value_but_fixme_should_propagate_errors();
+            paint(rect, *bitmap);
+            client().async_did_take_screenshot(m_id, bitmap->to_shareable_bitmap());
+        }
+    }
+
     if (!m_backing_stores.back_bitmap) {
         return;
     }
@@ -731,4 +751,10 @@ Web::PaintingCommandExecutorType PageClient::painting_command_executor_type() co
     return Web::PaintingCommandExecutorType::CPU;
 }
 
+void PageClient::queue_screenshot_task(Optional<i32> node_id)
+{
+    m_screenshot_tasks.enqueue({ node_id });
+    page().top_level_traversable()->set_needs_display();
+}
+
 }
diff --git a/Userland/Services/WebContent/PageClient.h b/Userland/Services/WebContent/PageClient.h
index 1cc06b4236f..063c836c94a 100644
--- a/Userland/Services/WebContent/PageClient.h
+++ b/Userland/Services/WebContent/PageClient.h
@@ -90,6 +90,8 @@ public:
 
     virtual Web::PaintingCommandExecutorType painting_command_executor_type() const override;
 
+    void queue_screenshot_task(Optional<i32> node_id);
+
 private:
     PageClient(PageHost&, u64 id);
 
@@ -186,6 +188,11 @@ private:
 
     PaintState m_paint_state { PaintState::Ready };
 
+    struct ScreenshotTask {
+        Optional<i32> node_id;
+    };
+    Queue<ScreenshotTask> m_screenshot_tasks;
+
     Web::CSS::PreferredColorScheme m_preferred_color_scheme { Web::CSS::PreferredColorScheme::Auto };
     Web::CSS::PreferredContrast m_preferred_contrast { Web::CSS::PreferredContrast::NoPreference };
     Web::CSS::PreferredMotion m_preferred_motion { Web::CSS::PreferredMotion::NoPreference };