From f03bebb9e31de206710ac29990184ac4b2108536 Mon Sep 17 00:00:00 2001 From: Your Name Date: Sat, 18 Jul 2026 20:18:28 +0530 Subject: [PATCH] [py][rb][bidi] Eliminate sleep-based polling for websocket command response --- .../webdriver/remote/websocket_connection.py | 40 ++++++++++--------- .../webdriver/common/websocket_connection.rb | 35 +++++++++++++++- 2 files changed, 54 insertions(+), 21 deletions(-) diff --git a/py/selenium/webdriver/remote/websocket_connection.py b/py/selenium/webdriver/remote/websocket_connection.py index 5d3f75f761542..9695b957e50ce 100644 --- a/py/selenium/webdriver/remote/websocket_connection.py +++ b/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) diff --git a/rb/lib/selenium/webdriver/common/websocket_connection.rb b/rb/lib/selenium/webdriver/common/websocket_connection.rb index 0774a46918009..2ede7f259bd5e 100644 --- a/rb/lib/selenium/webdriver/common/websocket_connection.rb +++ b/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