use aiohttp dns resolver - fixes #126

molotov/run.py

@@ -220,7 +220,8 @@ def main(args=None):
         if args.workers == 1:
             args.workers = 500
 
-    return run(args)
+    run(args)
+    return 0
 
 
 _SIZING = """\

molotov/session.py

@@ -66,13 +66,11 @@ def __init__(self, loop, console, verbose=0, statsd=None, resolve_dns=True, **kw
     async def send_event(self, event, **options):
         await self.eventer.send_event(event, session=self, **options)
 
-    def _dns_lookup(self, url):
-        return resolve(url)[0]
-
     async def _request(self, *args, **kw):
         args = list(args)
         if self._resolve_dns:
-            args[1] = self._dns_lookup(args[1])
+            resolved = await resolve(args[1])
+            args[1] = resolved[0]
         args = tuple(args)
         req = super(LoggedClientSession, self)._request
 

molotov/tests/test_util.py

@@ -4,6 +4,7 @@
 import unittest
 import os
 from molotov.util import resolve, expand_options, OptionError, set_var, get_var, _VARS
+from molotov.tests.support import async_test
 
 _HERE = os.path.dirname(__file__)
 config = os.path.join(_HERE, "..", "..", "molotov.json")
@@ -18,7 +19,8 @@ def setUp(self):
         super(TestUtil, self).setUp()
         _VARS.clear()
 
-    def test_resolve(self):
+    @async_test
+    async def test_resolve(self, loop, console, results):
 
         urls = [
             ("http://localhost:80/blah", "http://127.0.0.1:80/blah"),
@@ -33,7 +35,7 @@ def test_resolve(self):
         ]
 
         for url, wanted in urls:
-            changed, original, resolved = resolve(url)
+            changed, original, resolved = await resolve(url, loop=loop)
             self.assertEqual(changed, wanted, "%s vs %s" % (original, resolved))
 
     def test_config(self):

molotov/util.py

@@ -11,7 +11,9 @@
 import platform
 from urllib.parse import urlparse, urlunparse
 from socket import gethostbyname
+
 from aiohttp import ClientSession, __version__
+from aiohttp.resolver import DefaultResolver
 
 # this lib works for CPython 3.7+
 if platform.python_implementation() == "PyPy" or sys.version_info.minor < 7:
@@ -54,7 +56,15 @@ def is_stopped():
     return _STOP
 
 
-def resolve(url):
+_RESOLVERS = {}
+
+
+async def resolve(url, loop=None):
+    if loop in _RESOLVERS:
+        resolver = _RESOLVERS[loop]
+    else:
+        resolver = _RESOLVERS[loop] = DefaultResolver(loop=loop)
+
     parts = urlparse(url)
 
     if "@" in parts.netloc:
@@ -84,10 +94,13 @@ def resolve(url):
         resolved = _DNS_CACHE[host]
     else:
         try:
-            resolved = gethostbyname(host)
-            _DNS_CACHE[host] = resolved
+            hosts = await resolver.resolve(host, port)
         except socket.gaierror:
+            hosts = []
+        if len(hosts) == 0:
             return url, original, host
+        resolved = hosts[0]["host"]
+        _DNS_CACHE[host] = resolved
 
     # Don't use a resolved hostname for SSL requests otherwise the
     # certificate will not match the IP address (resolved)