From 375886bf95282ffc42c028c29b37ee0ac80bef75 Mon Sep 17 00:00:00 2001
From: YusukeIwaki <iwaki@i3-systems.com>
Date: Thu, 3 Dec 2020 03:55:14 +0900
Subject: [PATCH] Consider 'handle_message' is called just after 'send_text'
 before returning raw_send.

---
 lib/puppeteer/browser_runner.rb    |  1 -
 lib/puppeteer/cdp_session.rb       | 14 +++++++----
 lib/puppeteer/connection.rb        | 37 ++++++++++++++++++++++-------
 spec/puppeteer/cdp_session_spec.rb | 38 ++++++++++++++++++++++++++++++
 4 files changed, 75 insertions(+), 15 deletions(-)
 create mode 100644 spec/puppeteer/cdp_session_spec.rb

diff --git a/lib/puppeteer/browser_runner.rb b/lib/puppeteer/browser_runner.rb
index d6a8e528..842aab69 100644
--- a/lib/puppeteer/browser_runner.rb
+++ b/lib/puppeteer/browser_runner.rb
@@ -14,7 +14,6 @@ def initialize(executable_path, process_arguments, temp_directory)
     @proc = nil
     @connection = nil
     @closed = true
-    @listeners = []
   end
 
   attr_reader :proc, :connection
diff --git a/lib/puppeteer/cdp_session.rb b/lib/puppeteer/cdp_session.rb
index 28157c06..501f1586 100644
--- a/lib/puppeteer/cdp_session.rb
+++ b/lib/puppeteer/cdp_session.rb
@@ -9,7 +9,7 @@ class Error < StandardError; end
   # @param {string} targetType
   # @param {string} sessionId
   def initialize(connection, target_type, session_id)
-    @callbacks = {}
+    @callbacks = Concurrent::Hash.new
     @connection = connection
     @target_type = target_type
     @session_id = session_id
@@ -31,10 +31,14 @@ def async_send_message(method, params = {})
     if !@connection
       raise Error.new("Protocol error (#{method}): Session closed. Most likely the #{@target_type} has been closed.")
     end
-    id = @connection.raw_send(message: { sessionId: @session_id, method: method, params: params })
+
     promise = resolvable_future
-    callback = Puppeteer::Connection::MessageCallback.new(method: method, promise: promise)
-    @callbacks[id] = callback
+
+    @connection.generate_id do |id|
+      @callbacks[id] = Puppeteer::Connection::MessageCallback.new(method: method, promise: promise)
+      @connection.raw_send(id: id, message: { sessionId: @session_id, method: method, params: params })
+    end
+
     promise
   end
 
@@ -44,7 +48,7 @@ def handle_message(message)
       if callback = @callbacks.delete(message['id'])
         callback_with_message(callback, message)
       else
-        raise Error.new("unknown id: #{id}")
+        raise Error.new("unknown id: #{message['id']}")
       end
     else
       emit_event(message['method'], message['params'])
diff --git a/lib/puppeteer/connection.rb b/lib/puppeteer/connection.rb
index 63b8a5c1..86d27baf 100644
--- a/lib/puppeteer/connection.rb
+++ b/lib/puppeteer/connection.rb
@@ -39,7 +39,7 @@ def reject(error)
   def initialize(url, transport, delay = 0)
     @url = url
     @last_id = 0
-    @callbacks = {}
+    @callbacks = Concurrent::Hash.new
     @delay = delay
 
     @transport = transport
@@ -52,7 +52,7 @@ def initialize(url, transport, delay = 0)
       handle_close
     end
 
-    @sessions = {}
+    @sessions = Concurrent::Hash.new
     @closed = false
   end
 
@@ -92,22 +92,41 @@ def send_message(method, params = {})
   end
 
   def async_send_message(method, params = {})
-    id = raw_send(message: { method: method, params: params })
     promise = resolvable_future
-    @callbacks[id] = MessageCallback.new(method: method, promise: promise)
+
+    generate_id do |id|
+      @callbacks[id] = MessageCallback.new(method: method, promise: promise)
+      raw_send(id: id, message: { method: method, params: params })
+    end
+
     promise
   end
 
-  private def generate_id
-    @last_id += 1
+  # package private. not intended to use externally.
+  #
+  # ```usage
+  # connection.generate_id do |generated_id|
+  #   # play with generated_id
+  # end
+  # ````
+  #
+  def generate_id(&block)
+    block.call(@last_id += 1)
   end
 
-  def raw_send(message:)
-    id = generate_id
+  # package private. not intended to use externally.
+  def raw_send(id:, message:)
+    # In original puppeteer (JS) implementation,
+    # id is generated here using #generate_id and the id argument is not passed to #raw_send.
+    #
+    # However with concurrent-ruby, '#handle_message' is sometimes called
+    # just soon after @transport.send_text and **before returning the id.**
+    #
+    # So we have to know the message id in advance before send_text.
+    #
     payload = JSON.fast_generate(message.compact.merge(id: id))
     @transport.send_text(payload)
     request_debug_printer.handle_payload(payload)
-    id
   end
 
   # Just for effective debugging :)
diff --git a/spec/puppeteer/cdp_session_spec.rb b/spec/puppeteer/cdp_session_spec.rb
new file mode 100644
index 00000000..aa22a0c4
--- /dev/null
+++ b/spec/puppeteer/cdp_session_spec.rb
@@ -0,0 +1,38 @@
+require 'spec_helper'
+
+RSpec.describe Puppeteer::CDPSession do
+  let(:connection) { double(Puppeteer::Connection) }
+  let(:cdp_session_id) { SecureRandom.hex(16) }
+  let(:cdp_session) { Puppeteer::CDPSession.new(connection, 'page', cdp_session_id) }
+
+  describe '#send_message' do
+    before {
+      allow(connection).to receive(:generate_id) { |&block| block.call(SecureRandom.hex(16)) }
+      allow(connection).to receive(:raw_send) do |id:, message:|
+        Thread.new(id) do |message_id|
+          resp = {
+            'sessionId' => cdp_session_id,
+            'id' => message_id,
+            'result' => "pong",
+          }
+          cdp_session.handle_message(resp)
+        end
+      end
+    }
+
+    it 'should be thread safe' do
+      Timeout.timeout(5) do
+        await_all(1000.times.map { cdp_session.async_send_message('ping') })
+      end
+    end
+
+    it 'should raise error for unknown id' do
+      resp = {
+        'sessionId' => cdp_session_id,
+        'id' => -123,
+        'result' => "pong",
+      }
+      expect { cdp_session.handle_message(resp) }.to raise_error(/unknown id: -123/)
+    end
+  end
+end