From f23d00af07a0d3006cb5ea225e6a9fa8d2e2544d Mon Sep 17 00:00:00 2001
From: Kenta Murata <mrkn@mrkn.jp>
Date: Wed, 29 May 2024 12:48:32 +0900
Subject: [PATCH] Introduce PlaywrightManager helper module

---
 .../backend_helpers/playwright_manager.rb     | 71 +++++++++++++++++++
 lib/charty/backends/plotly.rb                 | 60 +++-------------
 2 files changed, 80 insertions(+), 51 deletions(-)
 create mode 100644 lib/charty/backends/backend_helpers/playwright_manager.rb

diff --git a/lib/charty/backends/backend_helpers/playwright_manager.rb b/lib/charty/backends/backend_helpers/playwright_manager.rb
new file mode 100644
index 00000000..84e695c9
--- /dev/null
+++ b/lib/charty/backends/backend_helpers/playwright_manager.rb
@@ -0,0 +1,71 @@
+module Charty
+  module Backends
+    module BackendHelpers
+      module PlaywrightManager
+        @_playwright_exec = nil
+        @_browser = nil
+        @_context = nil
+
+        at_exit do
+          if @_context
+            @_context.close
+            @_context = nil
+          end
+          if @_browser
+            @_browser.close
+            @_browser = nil
+          end
+          if @_playwright_exec
+            @_playwright_exec.stop
+            @_playwright_exec = nil
+          end
+        end
+
+        module_function def playwright
+          unless @_playwright_exec
+            load_playwright
+            path = ENV.fetch("PLAYWRIGHT_CLI_EXECUTABLE_PATH", "npx playwright")
+            @_playwright_exec = Playwright.create(playwright_cli_executable_path: path)
+          end
+          @_playwright_exec.playwright
+        end
+
+        module_function def launch_browser
+          playwright.chromium.launch(headless: true)
+        end
+
+        module_function def default_browser
+          unless @_browser
+            @_browser = launch_browser
+          end
+          @_browser
+        end
+
+        module_function def default_context
+          unless @_context
+            @_context = default_browser.new_context
+          end
+          @_context
+        end
+
+        module_function def new_page(&block)
+          page = default_context.new_page
+          return page unless block
+
+          begin
+            return block.call(page)
+          ensure
+            page.close
+          end
+        end
+
+        module_function def load_playwright
+          require "playwright"
+        rescue LoadError
+          $stderr.puts "ERROR: You need to install playwright and playwright-ruby-client before using Plotly renderer"
+          raise
+        end
+      end
+    end
+  end
+end
diff --git a/lib/charty/backends/plotly.rb b/lib/charty/backends/plotly.rb
index 06875031..1ad98738 100644
--- a/lib/charty/backends/plotly.rb
+++ b/lib/charty/backends/plotly.rb
@@ -5,6 +5,7 @@
 require_relative "plotly_helpers/html_renderer"
 require_relative "plotly_helpers/notebook_renderer"
 require_relative "plotly_helpers/plotly_renderer"
+require_relative "backend_helpers/playwright_manager"
 
 module Charty
   module Backends
@@ -918,60 +919,17 @@ def self.mathjax_config
         end
       end
 
-      @playwright_fiber = nil
-
-      def self.ensure_playwright
-        if @playwright_fiber.nil?
-          begin
-            require "playwright"
-          rescue LoadError
-            $stderr.puts "ERROR: You need to install playwright and playwright-ruby-client before using Plotly renderer"
-            raise
-          end
+      def self.render_image(input, output, format, element_id, width, height)
+        BackendHelpers::PlaywrightManager.new_page do |page|
+          page.set_viewport_size(width: width, height: height)
+          page.goto("file://#{input}")
+          element = page.query_selector("\##{element_id}")
 
-          @playwright_fiber = Fiber.new do
-            playwright_cli_executable_path = ENV.fetch("PLAYWRIGHT_CLI_EXECUTABLE_PATH", "npx playwright")
-            Playwright.create(playwright_cli_executable_path: playwright_cli_executable_path) do |playwright|
-              playwright.chromium.launch(headless: true) do |browser|
-                request = Fiber.yield
-                loop do
-                  result = nil
-                  case request.shift
-                  when :finish
-                    break
-                  when :render
-                    input, output, format, element_id, width, height = request
-
-                    page = browser.new_page
-                    page.set_viewport_size(width: width, height: height)
-                    page.goto("file://#{input}")
-                    element = page.query_selector("\##{element_id}")
-
-                    kwargs = {type: format}
-                    kwargs[:path] = output unless output.nil?
-                    result = element.screenshot(**kwargs)
-                  end
-                  request = Fiber.yield(result)
-                end
-              end
-            end
-          end
-          @playwright_fiber.resume
+          kwargs = {type: format}
+          kwargs[:path] = output unless output.nil?
+          element.screenshot(**kwargs)
         end
       end
-
-      def self.terminate_playwright
-        return if @playwright_fiber.nil?
-
-        @playwright_fiber.resume([:finish])
-      end
-
-      at_exit { terminate_playwright }
-
-      def self.render_image(input, output, format, element_id, width, height)
-        ensure_playwright if @playwright_fiber.nil?
-        @playwright_fiber.resume([:render, input, output, format.to_s, element_id, width, height])
-      end
     end
   end
 end