Skip to content

Commit

Permalink
[rh:websockets] Workaround race condition causing issues on PyPy (yt-…
Browse files Browse the repository at this point in the history
…dlp#9514)

Authored by: coletdjnz
  • Loading branch information
coletdjnz authored Mar 22, 2024
1 parent bc2b8c0 commit e5d4f11
Showing 1 changed file with 23 additions and 9 deletions.
32 changes: 23 additions & 9 deletions yt_dlp/networking/_websockets.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

import contextlib
import io
import logging
import ssl
Expand Down Expand Up @@ -38,27 +39,40 @@
import websockets.sync.client
from websockets.uri import parse_uri

# In websockets Connection, recv_exc and recv_events_exc are defined
# after the recv events handler thread is started [1].
# On our CI using PyPy, in some cases a race condition may occur
# where the recv events handler thread tries to use these attributes before they are defined [2].
# 1: https://github.com/python-websockets/websockets/blame/de768cf65e7e2b1a3b67854fb9e08816a5ff7050/src/websockets/sync/connection.py#L93
# 2: "AttributeError: 'ClientConnection' object has no attribute 'recv_events_exc'. Did you mean: 'recv_events'?"
import websockets.sync.connection # isort: split
with contextlib.suppress(Exception):
# > 12.0
websockets.sync.connection.Connection.recv_exc = None
# 12.0
websockets.sync.connection.Connection.recv_events_exc = None


class WebsocketsResponseAdapter(WebSocketResponse):

def __init__(self, wsw: websockets.sync.client.ClientConnection, url):
def __init__(self, ws: websockets.sync.client.ClientConnection, url):
super().__init__(
fp=io.BytesIO(wsw.response.body or b''),
fp=io.BytesIO(ws.response.body or b''),
url=url,
headers=wsw.response.headers,
status=wsw.response.status_code,
reason=wsw.response.reason_phrase,
headers=ws.response.headers,
status=ws.response.status_code,
reason=ws.response.reason_phrase,
)
self.wsw = wsw
self._ws = ws

def close(self):
self.wsw.close()
self._ws.close()
super().close()

def send(self, message):
# https://websockets.readthedocs.io/en/stable/reference/sync/client.html#websockets.sync.client.ClientConnection.send
try:
return self.wsw.send(message)
return self._ws.send(message)
except (websockets.exceptions.WebSocketException, RuntimeError, TimeoutError) as e:
raise TransportError(cause=e) from e
except SocksProxyError as e:
Expand All @@ -69,7 +83,7 @@ def send(self, message):
def recv(self):
# https://websockets.readthedocs.io/en/stable/reference/sync/client.html#websockets.sync.client.ClientConnection.recv
try:
return self.wsw.recv()
return self._ws.recv()
except SocksProxyError as e:
raise ProxyError(cause=e) from e
except (websockets.exceptions.WebSocketException, RuntimeError, TimeoutError) as e:
Expand Down

0 comments on commit e5d4f11

Please sign in to comment.