implement screenshot concurrency and its specs

YusukeIwaki · Dec 21, 2020 · 668ed36 · 668ed36
1 parent 1a49bfc
commit 668ed36
Show file tree

Hide file tree

Showing 14 changed files with 413 additions and 22 deletions.
diff --git a/lib/puppeteer/browser.rb b/lib/puppeteer/browser.rb
@@ -36,7 +36,6 @@ def initialize(connection:, context_ids:, ignore_https_errors:, default_viewport
     @ignore_https_errors = ignore_https_errors
     @default_viewport = default_viewport
     @process = process
-    # @screenshot_task_queue = TaskQueue.new
     @connection = connection
     @close_callback = close_callback
 
@@ -125,7 +124,6 @@ def handle_target_created(event)
       session_factory: -> { @connection.create_session(target_info) },
       ignore_https_errors: @ignore_https_errors,
       default_viewport: @default_viewport,
-      screenshot_task_queue: @screenshot_task_queue,
     )
     #   assert(!this._targets.has(event.targetInfo.targetId), 'Target should not exist before targetCreated');
     @targets[target_info.target_id] = target

diff --git a/lib/puppeteer/concurrent_ruby_utils.rb b/lib/puppeteer/concurrent_ruby_utils.rb
@@ -39,9 +39,9 @@ def await(future_or_value)
     end
   end
 
-  def future(&block)
-    Concurrent::Promises.future do
-      block.call
+  def future(*args, &block)
+    Concurrent::Promises.future(*args) do |*block_args|
+      block.call(*block_args)
     rescue => err
       Logger.new($stderr).warn(err)
       raise err

diff --git a/lib/puppeteer/page.rb b/lib/puppeteer/page.rb
@@ -3,6 +3,7 @@
 
 require_relative './page/pdf_options'
 require_relative './page/screenshot_options'
+require_relative './page/screenshot_task_queue'
 
 class Puppeteer::Page
   include Puppeteer::EventCallbackable
@@ -13,10 +14,9 @@ class Puppeteer::Page
   # @param {!Puppeteer.Target} target
   # @param {boolean} ignoreHTTPSErrors
   # @param {?Puppeteer.Viewport} defaultViewport
-  # @param {!Puppeteer.TaskQueue} screenshotTaskQueue
   # @return {!Promise<!Page>}
-  def self.create(client, target, ignore_https_errors, default_viewport, screenshot_task_queue)
-    page = Puppeteer::Page.new(client, target, ignore_https_errors, screenshot_task_queue)
+  def self.create(client, target, ignore_https_errors, default_viewport)
+    page = Puppeteer::Page.new(client, target, ignore_https_errors)
     page.init
     if default_viewport
       page.viewport = default_viewport
@@ -27,8 +27,7 @@ def self.create(client, target, ignore_https_errors, default_viewport, screensho
   # @param {!Puppeteer.CDPSession} client
   # @param {!Puppeteer.Target} target
   # @param {boolean} ignoreHTTPSErrors
-  # @param {!Puppeteer.TaskQueue} screenshotTaskQueue
-  def initialize(client, target, ignore_https_errors, screenshot_task_queue)
+  def initialize(client, target, ignore_https_errors)
     @closed = false
     @client = client
     @target = target
@@ -43,7 +42,7 @@ def initialize(client, target, ignore_https_errors, screenshot_task_queue)
     @page_bindings = {}
     # @coverage = Coverage.new(client)
     @javascript_enabled = true
-    @screenshot_task_queue = screenshot_task_queue
+    @screenshot_task_queue = ScreenshotTaskQueue.new
 
     @workers = {}
     @client.on_event('Target.attachedToTarget') do |event|
@@ -867,15 +866,28 @@ def title
     main_frame.title
   end
 
-  # /**
-  #  * @param {!ScreenshotOptions=} options
-  #  * @return {!Promise<!Buffer|!String>}
-  #  */
-  def screenshot(options = {})
+  # @param type [String] "png"|"jpeg"
+  # @param path [String]
+  # @param full_page [Boolean]
+  # @param clip [Hash]
+  # @param quality [Integer]
+  # @param omit_background [Boolean]
+  # @param encoding [String]
+  def screenshot(type: nil, path: nil, full_page: nil, clip: nil, quality: nil, omit_background: nil, encoding: nil)
+    options = {
+      type: type,
+      path: path,
+      full_page: full_page,
+      clip: clip,
+      quality:  quality,
+      omit_background: omit_background,
+      encoding: encoding,
+    }.compact
     screenshot_options = ScreenshotOptions.new(options)
 
-    # @screenshot_task_queue.post_task(-> { screenshot_task(screenshot_options.type, screenshot_options) })
-    screenshot_task(screenshot_options.type, screenshot_options)
+    @screenshot_task_queue.post_task do
+      screenshot_task(screenshot_options.type, screenshot_options)
+    end
   end
 
   # @param {"png"|"jpeg"} format

diff --git a/lib/puppeteer/page/screenshot_task_queue.rb b/lib/puppeteer/page/screenshot_task_queue.rb
@@ -0,0 +1,13 @@
+class Puppeteer::Page
+  class ScreenshotTaskQueue
+    def initialize
+      @chain = Concurrent::Promises.fulfilled_future(nil)
+    end
+
+    def post_task(&block)
+      result = @chain.then { block.call }
+      @chain = result.rescue { nil }
+      result.value!
+    end
+  end
+end
diff --git a/lib/puppeteer/target.rb b/lib/puppeteer/target.rb
@@ -18,15 +18,13 @@ def initialize(options)
   # @param {!function():!Promise<!Puppeteer.CDPSession>} sessionFactory
   # @param {boolean} ignoreHTTPSErrors
   # @param {?Puppeteer.Viewport} defaultViewport
-  # @param {!Puppeteer.TaskQueue} screenshotTaskQueue
-  def initialize(target_info:, browser_context:, session_factory:, ignore_https_errors:, default_viewport:, screenshot_task_queue:)
+  def initialize(target_info:, browser_context:, session_factory:, ignore_https_errors:, default_viewport:)
     @target_info = target_info
     @browser_context = browser_context
     @target_id = target_info.target_id
     @session_factory = session_factory
     @ignore_https_errors = ignore_https_errors
     @default_viewport = default_viewport
-    @screenshot_task_queue = screenshot_task_queue
 
 
     #    /** @type {?Promise<!Puppeteer.Page>} */
@@ -87,7 +85,7 @@ def create_cdp_session
   def page
     if ['page', 'background_page', 'webview'].include?(@target_info.type) && @page.nil?
       client = @session_factory.call
-      @page = Puppeteer::Page.create(client, self, @ignore_https_errors, @default_viewport, @screenshot_task_queue)
+      @page = Puppeteer::Page.create(client, self, @ignore_https_errors, @default_viewport)
     end
     @page
   end

diff --git a/puppeteer-ruby.gemspec b/puppeteer-ruby.gemspec
@@ -22,6 +22,7 @@ Gem::Specification.new do |spec|
   spec.add_dependency 'websocket-driver', '>= 0.6.0'
   spec.add_dependency 'mime-types', '>= 3.0'
   spec.add_development_dependency 'bundler', '~> 1.17'
+  spec.add_development_dependency 'chunky_png'
   spec.add_development_dependency 'pry-byebug'
   spec.add_development_dependency 'rake', '~> 10.0'
   spec.add_development_dependency 'rspec', '~> 3.0'

diff --git a/spec/integration/golden-chromium/grid-cell-1.png b/spec/integration/golden-chromium/grid-cell-1.png
diff --git a/spec/integration/golden-chromium/screenshot-clip-rect.png b/spec/integration/golden-chromium/screenshot-clip-rect.png
diff --git a/spec/integration/golden-chromium/screenshot-grid-fullpage.png b/spec/integration/golden-chromium/screenshot-grid-fullpage.png
diff --git a/spec/integration/golden-chromium/screenshot-offscreen-clip.png b/spec/integration/golden-chromium/screenshot-offscreen-clip.png
diff --git a/spec/integration/golden-chromium/screenshot-sanity.png b/spec/integration/golden-chromium/screenshot-sanity.png