1  from __future__ import annotations
   2  
   3  import asyncio
   4  import socket
   5  import threading
   6  from collections.abc import Callable
   7  from selectors import EVENT_READ, EVENT_WRITE, DefaultSelector
   8  from typing import TYPE_CHECKING, Any
   9  
  10  if TYPE_CHECKING:
  11      from _typeshed import FileDescriptorLike
  12  
  13  _selector_lock = threading.Lock()
  14  _selector: Selector | None = None
  15  
  16  
  17  class Selector:
  18      def __init__(self) -> None:
  19          self._thread = threading.Thread(target=self.run, name="AnyIO socket selector")
  20          self._selector = DefaultSelector()
  21          self._send, self._receive = socket.socketpair()
  22          self._send.setblocking(False)
  23          self._receive.setblocking(False)
  24          # This somewhat reduces the amount of memory wasted queueing up data
  25          # for wakeups. With these settings, maximum number of 1-byte sends
  26          # before getting BlockingIOError:
  27          #   Linux 4.8: 6
  28          #   macOS (darwin 15.5): 1
  29          #   Windows 10: 525347
  30          # Windows you're weird. (And on Windows setting SNDBUF to 0 makes send
  31          # blocking, even on non-blocking sockets, so don't do that.)
  32          self._receive.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 1)
  33          self._send.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1)
  34          # On Windows this is a TCP socket so this might matter. On other
  35          # platforms this fails b/c AF_UNIX sockets aren't actually TCP.
  36          try:
  37              self._send.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
  38          except OSError:
  39              pass
  40  
  41          self._selector.register(self._receive, EVENT_READ)
  42          self._closed = False
  43  
  44      def start(self) -> None:
  45          self._thread.start()
  46          threading._register_atexit(self._stop)  # type: ignore[attr-defined]
  47  
  48      def _stop(self) -> None:
  49          global _selector
  50          self._closed = True
  51          self._notify_self()
  52          self._send.close()
  53          self._thread.join()
  54          self._selector.unregister(self._receive)
  55          self._receive.close()
  56          self._selector.close()
  57          _selector = None
  58          assert not self._selector.get_map(), (
  59              "selector still has registered file descriptors after shutdown"
  60          )
  61  
  62      def _notify_self(self) -> None:
  63          try:
  64              self._send.send(b"\x00")
  65          except BlockingIOError:
  66              pass
  67  
  68      def add_reader(self, fd: FileDescriptorLike, callback: Callable[[], Any]) -> None:
  69          loop = asyncio.get_running_loop()
  70          try:
  71              key = self._selector.get_key(fd)
  72          except KeyError:
  73              self._selector.register(fd, EVENT_READ, {EVENT_READ: (loop, callback)})
  74          else:
  75              if EVENT_READ in key.data:
  76                  raise ValueError(
  77                      "this file descriptor is already registered for reading"
  78                  )
  79  
  80              key.data[EVENT_READ] = loop, callback
  81              self._selector.modify(fd, key.events | EVENT_READ, key.data)
  82  
  83          self._notify_self()
  84  
  85      def add_writer(self, fd: FileDescriptorLike, callback: Callable[[], Any]) -> None:
  86          loop = asyncio.get_running_loop()
  87          try:
  88              key = self._selector.get_key(fd)
  89          except KeyError:
  90              self._selector.register(fd, EVENT_WRITE, {EVENT_WRITE: (loop, callback)})
  91          else:
  92              if EVENT_WRITE in key.data:
  93                  raise ValueError(
  94                      "this file descriptor is already registered for writing"
  95                  )
  96  
  97              key.data[EVENT_WRITE] = loop, callback
  98              self._selector.modify(fd, key.events | EVENT_WRITE, key.data)
  99  
 100          self._notify_self()
 101  
 102      def remove_reader(self, fd: FileDescriptorLike) -> bool:
 103          try:
 104              key = self._selector.get_key(fd)
 105          except KeyError:
 106              return False
 107  
 108          if new_events := key.events ^ EVENT_READ:
 109              del key.data[EVENT_READ]
 110              self._selector.modify(fd, new_events, key.data)
 111          else:
 112              self._selector.unregister(fd)
 113  
 114          return True
 115  
 116      def remove_writer(self, fd: FileDescriptorLike) -> bool:
 117          try:
 118              key = self._selector.get_key(fd)
 119          except KeyError:
 120              return False
 121  
 122          if new_events := key.events ^ EVENT_WRITE:
 123              del key.data[EVENT_WRITE]
 124              self._selector.modify(fd, new_events, key.data)
 125          else:
 126              self._selector.unregister(fd)
 127  
 128          return True
 129  
 130      def run(self) -> None:
 131          while not self._closed:
 132              for key, events in self._selector.select():
 133                  if key.fileobj is self._receive:
 134                      try:
 135                          while self._receive.recv(4096):
 136                              pass
 137                      except BlockingIOError:
 138                          pass
 139  
 140                      continue
 141  
 142                  if events & EVENT_READ:
 143                      loop, callback = key.data[EVENT_READ]
 144                      self.remove_reader(key.fd)
 145                      try:
 146                          loop.call_soon_threadsafe(callback)
 147                      except RuntimeError:
 148                          pass  # the loop was already closed
 149  
 150                  if events & EVENT_WRITE:
 151                      loop, callback = key.data[EVENT_WRITE]
 152                      self.remove_writer(key.fd)
 153                      try:
 154                          loop.call_soon_threadsafe(callback)
 155                      except RuntimeError:
 156                          pass  # the loop was already closed
 157  
 158  
 159  def get_selector() -> Selector:
 160      global _selector
 161  
 162      with _selector_lock:
 163          if _selector is None:
 164              _selector = Selector()
 165              _selector.start()
 166  
 167          return _selector
 168