Skip to content

Commit b2e3c9d

Browse files
authored
Migrate to using efficient I/O multiplexing selectors (#865)
* Migrate to using efficient I/O multiplexing selectors * naming
1 parent 031d38c commit b2e3c9d

File tree

1 file changed

+23
-23
lines changed

1 file changed

+23
-23
lines changed

pychromecast/socket_client.py

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -10,31 +10,31 @@
1010
from __future__ import annotations
1111

1212
import abc
13-
from dataclasses import dataclass
1413
import errno
1514
import json
1615
import logging
17-
import select
16+
import selectors
1817
import socket
1918
import ssl
2019
import threading
2120
import time
2221
from collections import defaultdict
22+
from dataclasses import dataclass
2323
from struct import pack, unpack
2424

2525
import zeroconf
2626

27-
from .controllers import CallbackType, BaseController
27+
from .const import MESSAGE_TYPE, REQUEST_ID, SESSION_ID
28+
from .controllers import BaseController, CallbackType
2829
from .controllers.media import MediaController
2930
from .controllers.receiver import CastStatus, CastStatusListener, ReceiverController
30-
from .const import MESSAGE_TYPE, REQUEST_ID, SESSION_ID
3131
from .dial import get_host_from_service
3232
from .error import (
3333
ChromecastConnectionError,
3434
ControllerNotRegistered,
35-
UnsupportedNamespace,
3635
NotConnected,
3736
PyChromecastStopped,
37+
UnsupportedNamespace,
3838
)
3939

4040
# pylint: disable-next=no-name-in-module
@@ -64,8 +64,6 @@
6464
CONNECTION_STATUS_FAILED_RESOLVE = "FAILED_RESOLVE"
6565
# The socket connection was lost and needs to be retried
6666
CONNECTION_STATUS_LOST = "LOST"
67-
# Check for select poll method
68-
SELECT_HAS_POLL = hasattr(select, "poll")
6967

7068
HB_PING_TIME = 10
7169
HB_PONG_TIME = 10
@@ -213,6 +211,11 @@ def __init__(
213211
self.connecting = True
214212
self.first_connection = True
215213
self.socket: socket.socket | ssl.SSLSocket | None = None
214+
self.selector = selectors.DefaultSelector()
215+
self.wakeup_selector_key = self.selector.register(
216+
self.socketpair[0], selectors.EVENT_READ
217+
)
218+
self.remote_selector_key: selectors.SelectorKey | None = None
216219

217220
# dict mapping namespace on Controller objects
218221
self._handlers: dict[str, set[BaseController]] = defaultdict(set)
@@ -236,8 +239,10 @@ def initialize_connection( # pylint:disable=too-many-statements, too-many-branc
236239
tries = self.tries
237240

238241
if self.socket is not None:
242+
self.selector.unregister(self.socket)
239243
self.socket.close()
240244
self.socket = None
245+
self.remote_selector_key = None
241246

242247
# Make sure nobody is blocking.
243248
for callback_function in self._request_callbacks.values():
@@ -286,10 +291,15 @@ def mdns_backoff(
286291
try:
287292
if self.socket is not None:
288293
# If we retry connecting, we need to clean up the socket again
289-
self.socket.close() # type: ignore[unreachable]
294+
self.selector.unregister(self.socket) # type: ignore[unreachable]
295+
self.socket.close()
290296
self.socket = None
297+
self.remote_selector_key = None
291298

292299
self.socket = new_socket()
300+
self.remote_selector_key = self.selector.register(
301+
self.socket, selectors.EVENT_READ
302+
)
293303
self.socket.settimeout(self.timeout)
294304
self._report_connection_status(
295305
ConnectionStatus(
@@ -557,20 +567,8 @@ def _run_once(self) -> int:
557567
assert self.socket is not None
558568

559569
# poll the socket, as well as the socketpair to allow us to be interrupted
560-
rlist = [self.socket, self.socketpair[0]]
561570
try:
562-
if SELECT_HAS_POLL is True:
563-
# Map file descriptors to socket objects because select.select does not support fd > 1024
564-
# https://stackoverflow.com/questions/14250751/how-to-increase-filedescriptors-range-in-python-select
565-
fd_to_socket = {rlist_item.fileno(): rlist_item for rlist_item in rlist}
566-
567-
poll_obj = select.poll()
568-
for poll_fd in rlist:
569-
poll_obj.register(poll_fd, select.POLLIN)
570-
poll_result = poll_obj.poll()
571-
can_read = [fd_to_socket[fd] for fd, _status in poll_result]
572-
else:
573-
can_read, _, _ = select.select(rlist, [], [], None)
571+
ready = self.selector.select()
574572
except (ValueError, OSError) as exc:
575573
self.logger.error(
576574
"[%s(%s):%s] Error in select call: %s",
@@ -582,9 +580,10 @@ def _run_once(self) -> int:
582580
self._force_recon = True
583581
return 0
584582

583+
can_read = {key for key, _ in ready}
585584
# read message from chromecast
586585
message = None
587-
if self.socket in can_read and not self._force_recon:
586+
if self.remote_selector_key in can_read and not self._force_recon:
588587
try:
589588
message = self._read_message()
590589
except InterruptLoop as exc:
@@ -620,7 +619,7 @@ def _run_once(self) -> int:
620619
else:
621620
data = _dict_from_message_payload(message)
622621

623-
if self.socketpair[0] in can_read:
622+
if self.wakeup_selector_key in can_read:
624623
# Clear the socket's buffer
625624
self.socketpair[0].recv(128)
626625

@@ -765,6 +764,7 @@ def _cleanup(self) -> None:
765764

766765
self.socketpair[0].close()
767766
self.socketpair[1].close()
767+
self.selector.close()
768768

769769
self.connecting = True
770770

0 commit comments

Comments
 (0)