Use c-ares event thread when available (#145)

Author: bdraco

README.rst

@@ -58,9 +58,17 @@ The API is pretty simple, three functions are provided in the ``DNSResolver`` cl
 Note for Windows users
 ======================
 
-This library requires the asyncio loop to be a `SelectorEventLoop`, which is not the default on Windows.
+This library requires the use of an ``asyncio.SelectorEventLoop`` on Windows
+**only** when using a custom build of ``pycares`` that links against a system-
+provided ``c-ares`` library **without** thread-safety support. This is because
+non-thread-safe builds of ``c-ares`` are incompatible with the default
+``ProactorEventLoop`` on Windows.
 
-The default can be changed as follows (do this very early in your application):
+If you're using the official prebuilt ``pycares`` wheels on PyPI (version 4.7.0 or
+later), which include a thread-safe version of ``c-ares``, this limitation does
+**not** apply and can be safely ignored.
+
+The default event loop can be changed as follows (do this very early in your application):
 
 .. code:: python
 

aiodns/__init__.py

@@ -7,6 +7,7 @@
 
 from typing import (
     Any,
+    Callable,
     Optional,
     Set,
     Sequence,
@@ -21,6 +22,11 @@
 
 __all__ = ('DNSResolver', 'error')
 
+WINDOWS_SELECTOR_ERR_MSG = (
+    "aiodns needs a SelectorEventLoop on Windows. See more: "
+    "https://github.com/aio-libs/aiodns#note-for-windows-users"
+)
+
 
 READ = 1
 WRITE = 2
@@ -52,22 +58,26 @@ def __init__(self, nameservers: Optional[Sequence[str]] = None,
                  **kwargs: Any) -> None:
         self.loop = loop or asyncio.get_event_loop()
         assert self.loop is not None
-        if sys.platform == 'win32':
-            if not isinstance(self.loop, asyncio.SelectorEventLoop):
+        kwargs.pop('sock_state_cb', None)
+        timeout = kwargs.pop('timeout', None)
+        self._timeout = timeout
+        self._event_thread = hasattr(pycares,"ares_threadsafety") and pycares.ares_threadsafety()
+        if self._event_thread:
+            # pycares is thread safe
+            self._channel = pycares.Channel(event_thread=True,
+                                            timeout=timeout,
+                                            **kwargs)
+        else:
+            if sys.platform == 'win32' and not isinstance(self.loop, asyncio.SelectorEventLoop):
                 try:
                     import winloop
                     if not isinstance(self.loop , winloop.Loop):
-                        raise RuntimeError(
-                            'aiodns needs a SelectorEventLoop on Windows. See more: https://github.com/saghul/aiodns/issues/86')
+                        raise RuntimeError(WINDOWS_SELECTOR_ERR_MSG)
                 except ModuleNotFoundError:
-                    raise RuntimeError(
-                        'aiodns needs a SelectorEventLoop on Windows. See more: https://github.com/saghul/aiodns/issues/86')
-        kwargs.pop('sock_state_cb', None)
-        timeout = kwargs.pop('timeout', None)
-        self._timeout = timeout
-        self._channel = pycares.Channel(sock_state_cb=self._sock_state_cb,
-                                        timeout=timeout,
-                                        **kwargs)
+                    raise RuntimeError(WINDOWS_SELECTOR_ERR_MSG)
+            self._channel = pycares.Channel(sock_state_cb=self._sock_state_cb,
+                                            timeout=timeout,
+                                            **kwargs)
         if nameservers:
             self.nameservers = nameservers
         self._read_fds = set() # type: Set[int]
@@ -91,6 +101,20 @@ def _callback(fut: asyncio.Future, result: Any, errorno: int) -> None:
         else:
             fut.set_result(result)
 
+    def _get_future_callback(self) -> Tuple["asyncio.Future[Any]", Callable[[Any, int], None]]:
+        """Return a future and a callback to set the result of the future."""
+        cb: Callable[[Any, int], None]
+        future: "asyncio.Future[Any]" = self.loop.create_future()
+        if self._event_thread:
+            cb = functools.partial(  # type: ignore[assignment]
+                self.loop.call_soon_threadsafe,
+                self._callback,  # type: ignore[arg-type]
+                future
+            )
+        else:
+            cb = functools.partial(self._callback, future)
+        return future, cb
+
     def query(self, host: str, qtype: str, qclass: Optional[str]=None) -> asyncio.Future:
         try:
             qtype = query_type_map[qtype]
@@ -102,32 +126,27 @@ def query(self, host: str, qtype: str, qclass: Optional[str]=None) -> asyncio.Fu
             except KeyError:
                 raise ValueError('invalid query class: {}'.format(qclass))
 
-        fut = asyncio.Future(loop=self.loop)  # type: asyncio.Future
-        cb = functools.partial(self._callback, fut)
+        fut, cb = self._get_future_callback()
         self._channel.query(host, qtype, cb, query_class=qclass)
         return fut
 
     def gethostbyname(self, host: str, family: socket.AddressFamily) -> asyncio.Future:
-        fut = asyncio.Future(loop=self.loop)  # type: asyncio.Future
-        cb = functools.partial(self._callback, fut)
+        fut, cb = self._get_future_callback()
         self._channel.gethostbyname(host, family, cb)
         return fut
 
     def getaddrinfo(self, host: str, family: socket.AddressFamily = socket.AF_UNSPEC, port: Optional[int] = None, proto: int = 0, type: int = 0, flags: int = 0) -> asyncio.Future:
-        fut = asyncio.Future(loop=self.loop)  # type: asyncio.Future
-        cb = functools.partial(self._callback, fut)
+        fut, cb = self._get_future_callback()
         self._channel.getaddrinfo(host, port, cb, family=family, type=type, proto=proto, flags=flags)
         return fut
 
     def getnameinfo(self, sockaddr: Union[Tuple[str, int], Tuple[str, int, int, int]], flags: int = 0) -> asyncio.Future:
-        fut = asyncio.Future(loop=self.loop)  # type: asyncio.Future
-        cb = functools.partial(self._callback, fut)
+        fut, cb = self._get_future_callback()
         self._channel.getnameinfo(sockaddr, flags, cb)
         return fut
 
     def gethostbyaddr(self, name: str) -> asyncio.Future:
-        fut = asyncio.Future(loop=self.loop)  # type: asyncio.Future
-        cb = functools.partial(self._callback, fut)
+        fut, cb = self._get_future_callback()
         self._channel.gethostbyaddr(name, cb)
         return fut
 

tests/test_aiodns.py

@@ -3,9 +3,11 @@
 import asyncio
 import ipaddress
 import unittest
+import pytest
 import socket
 import sys
 import time
+import unittest.mock
 
 import aiodns
 
@@ -211,5 +213,25 @@ def setUp(self):
         self.resolver = aiodns.DNSResolver(loop=self.loop, timeout=5.0)
         self.resolver.nameservers = ['8.8.8.8']
 
+
+class TestNoEventThreadDNS(DNSTest):
+    """Test DNSResolver with no event thread."""
+
+    def setUp(self):
+        with unittest.mock.patch('aiodns.pycares.ares_threadsafety', return_value=False):
+            super().setUp()
+
+
+@unittest.skipIf(sys.platform != 'win32', 'Only run on Windows')
+def test_win32_no_selector_event_loop():
+    """Test DNSResolver with Windows without SelectorEventLoop."""
+    asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy())
+    with (
+        pytest.raises(RuntimeError, match="aiodns needs a SelectorEventLoop on Windows"),
+        unittest.mock.patch('aiodns.pycares.ares_threadsafety', return_value=False)
+    ):
+        aiodns.DNSResolver(loop=asyncio.new_event_loop(), timeout=5.0)
+
+
 if __name__ == "__main__":  # pragma: no cover
     unittest.main(verbosity=2)