[py][rb][bidi] Eliminate sleep-based polling for websocket command response

py/selenium/webdriver/remote/websocket_connection.py

@@ -96,10 +96,13 @@ def __init__(self, url, timeout, interval):
         self._id = 0
         self._id_lock = threading.Lock()
         self._messages = {}
-        self._started = False
+        self._started = threading.Event()
+        self._events = {}
+        self._events_lock = threading.Lock()
 
         self._start_ws()
-        self._wait_until(lambda: self._started)
+        if not self._started.wait(timeout=self.response_wait_timeout):
+            raise WebDriverException("Timed out waiting for connection to start")
 
     def close(self):
         # Close the socket first so ``run_forever`` returns; only then join the
@@ -112,13 +115,18 @@ def close(self):
                 logger.debug(f"Error while closing websocket connection: {e}")
         if self._ws_thread is not None:
             self._ws_thread.join(timeout=self.response_wait_timeout)
-        self._started = False
+        self._started.clear()
         self._ws = None
 
     def execute(self, command):
         with self._id_lock:
             self._id += 1
             current_id = self._id
+
+        event = threading.Event()
+        with self._events_lock:
+            self._events[current_id] = event
+
         payload = self._serialize_command(command)
         payload["id"] = current_id
         if self.session_id:
@@ -128,8 +136,9 @@ def execute(self, command):
         logger.debug(f"-> {data}"[: self._max_log_message_size])
         self._ws.send(data)
 
-        self._wait_until(lambda: current_id in self._messages)
-        if current_id not in self._messages:
+        if not event.wait(timeout=self.response_wait_timeout):
+            with self._events_lock:
+                self._events.pop(current_id, None)
             raise WebDriverException(f"Timed out waiting for response to BiDi command {current_id}")
         response = self._messages.pop(current_id)
 
@@ -177,7 +186,7 @@ def _deserialize_result(self, result, command):
 
     def _start_ws(self):
         def on_open(ws):
-            self._started = True
+            self._started.set()
 
         def on_message(ws, message):
             self._process_message(message)
@@ -201,21 +210,14 @@ def _process_message(self, message):
         logger.debug(f"<- {message}"[: self._max_log_message_size])
 
         if "id" in message:
-            self._messages[message["id"]] = message
+            msg_id = message["id"]
+            self._messages[msg_id] = message
+            with self._events_lock:
+                event = self._events.pop(msg_id, None)
+            if event:
+                event.set()
 
         if "method" in message:
             params = message["params"]
             for callback in self.callbacks.get(message["method"], []):
                 Thread(target=callback, args=(params,), daemon=True).start()
-
-    def _wait_until(self, condition):
-        timeout = self.response_wait_timeout
-        interval = self.response_wait_interval
-
-        while timeout > 0:
-            result = condition()
-            if result:
-                return result
-            else:
-                timeout -= interval
-                sleep(interval)

rb/lib/selenium/webdriver/common/websocket_connection.rb

@@ -45,6 +45,7 @@ def initialize(url:)
         @closing = false
         @session_id = nil
         @url = url
+        @pending_responses = {}
 
         process_handshake
         @socket_thread = attach_socket_listener
@@ -103,13 +104,33 @@ def send_cmd(**payload)
         data = JSON.generate(data)
         out_frame = WebSocket::Frame::Outgoing::Client.new(version: ws.version, data: data, type: 'text')
 
+        cond = ConditionVariable.new
+        @messages_mtx.synchronize do
+          @pending_responses[id] = { cond: cond, response: nil }
+        end
+
         begin
           socket.write(out_frame.to_s)
         rescue *CONNECTION_ERRORS => e
+          @messages_mtx.synchronize { @pending_responses.delete(id) }
           raise e, "WebSocket is closed (#{e.class}: #{e.message})"
         end
 
-        wait.until { @messages_mtx.synchronize { messages.delete(id) } }
+        @messages_mtx.synchronize do
+          timeout_at = Process.clock_gettime(Process::CLOCK_MONOTONIC) + RESPONSE_WAIT_TIMEOUT
+          while @pending_responses[id] && @pending_responses[id][:response].nil?
+            remaining = timeout_at - Process.clock_gettime(Process::CLOCK_MONOTONIC)
+            if remaining <= 0
+              @pending_responses.delete(id)
+              raise Error::TimeoutError, "timed out after #{RESPONSE_WAIT_TIMEOUT} seconds waiting for response to command #{id}"
+            end
+            cond.wait(@messages_mtx, remaining)
+          end
+          entry = @pending_responses.delete(id)
+          raise Error::TimeoutError, "timed out after #{RESPONSE_WAIT_TIMEOUT} seconds waiting for response to command #{id}" unless entry && entry[:response]
+
+          entry[:response]
+        end
       end
 
       private
@@ -159,7 +180,17 @@ def process_frame(frame)
         return {} if message.empty?
 
         msg = JSON.parse(message)
-        @messages_mtx.synchronize { messages[msg['id']] = msg if msg.key?('id') }
+        if msg.key?('id')
+          msg_id = msg['id']
+          @messages_mtx.synchronize do
+            if @pending_responses.key?(msg_id)
+              @pending_responses[msg_id][:response] = msg
+              @pending_responses[msg_id][:cond].signal
+            else
+              messages[msg_id] = msg
+            end
+          end
+        end
 
         WebDriver.logger.debug "WebSocket <- #{msg}"[...MAX_LOG_MESSAGE_SIZE], id: :ws
         msg