[PR #12106/8ab84c52 backport][3.14] Bound DNS cache (#12116)

**This is a backport of PR #12106 as merged into master
(8ab84c5).**

---------

Co-authored-by: Sam Bull <[email protected]>
Co-authored-by: gonas <[email protected]>

Author: 3 people

CHANGES/12106.feature.rst

@@ -0,0 +1 @@
+Added a ``dns_cache_max_size`` parameter to ``TCPConnector`` to limit the size of the cache -- by :user:`Dreamsorcerer`.

aiohttp/connector.py

@@ -839,25 +839,33 @@ async def _create_connection(
 
 
 class _DNSCacheTable:
-    def __init__(self, ttl: float | None = None) -> None:
-        self._addrs_rr: dict[tuple[str, int], tuple[Iterator[ResolveResult], int]] = {}
+    def __init__(self, ttl: float | None = None, max_size: int = 1000) -> None:
+        self._addrs_rr: OrderedDict[
+            tuple[str, int], tuple[Iterator[ResolveResult], int]
+        ] = OrderedDict()
         self._timestamps: dict[tuple[str, int], float] = {}
         self._ttl = ttl
+        self._max_size = max_size
 
     def __contains__(self, host: object) -> bool:
         return host in self._addrs_rr
 
     def add(self, key: tuple[str, int], addrs: list[ResolveResult]) -> None:
+        if key in self._addrs_rr:
+            self._addrs_rr.move_to_end(key)
+
         self._addrs_rr[key] = (cycle(addrs), len(addrs))
 
         if self._ttl is not None:
             self._timestamps[key] = monotonic()
 
+        if len(self._addrs_rr) > self._max_size:
+            oldest_key, _ = self._addrs_rr.popitem(last=False)
+            self._timestamps.pop(oldest_key, None)
+
     def remove(self, key: tuple[str, int]) -> None:
         self._addrs_rr.pop(key, None)
-
-        if self._ttl is not None:
-            self._timestamps.pop(key, None)
+        self._timestamps.pop(key, None)
 
     def clear(self) -> None:
         self._addrs_rr.clear()
@@ -868,6 +876,7 @@ def next_addrs(self, key: tuple[str, int]) -> list[ResolveResult]:
         addrs = list(islice(loop, length))
         # Consume one more element to shift internal state of `cycle`
         next(loop)
+        self._addrs_rr.move_to_end(key)
         return addrs
 
     def expired(self, key: tuple[str, int]) -> bool:
@@ -956,6 +965,7 @@ def __init__(
         fingerprint: bytes | None = None,
         use_dns_cache: bool = True,
         ttl_dns_cache: int | None = 10,
+        dns_cache_max_size: int = 1000,
         family: socket.AddressFamily = socket.AddressFamily.AF_UNSPEC,
         ssl_context: SSLContext | None = None,
         ssl: bool | Fingerprint | SSLContext = True,
@@ -994,7 +1004,9 @@ def __init__(
             self._resolver_owner = False
 
         self._use_dns_cache = use_dns_cache
-        self._cached_hosts = _DNSCacheTable(ttl=ttl_dns_cache)
+        self._cached_hosts = _DNSCacheTable(
+            ttl=ttl_dns_cache, max_size=dns_cache_max_size
+        )
         self._throttle_dns_futures: dict[tuple[str, int], set[asyncio.Future[None]]] = (
             {}
         )

tests/test_connector.py

@@ -4027,6 +4027,25 @@ async def handler(request):
 
 
 class TestDNSCacheTable:
+    host1 = ("localhost", 80)
+    host2 = ("foo", 80)
+    result1: ResolveResult = {
+        "hostname": "localhost",
+        "host": "127.0.0.1",
+        "port": 80,
+        "family": socket.AF_INET,
+        "proto": 0,
+        "flags": socket.AI_NUMERICHOST,
+    }
+    result2: ResolveResult = {
+        "hostname": "foo",
+        "host": "127.0.0.2",
+        "port": 80,
+        "family": socket.AF_INET,
+        "proto": 0,
+        "flags": socket.AI_NUMERICHOST,
+    }
+
     @pytest.fixture
     def dns_cache_table(self):
         return _DNSCacheTable()
@@ -4112,6 +4131,63 @@ def test_next_addrs_single(self, dns_cache_table) -> None:
         addrs = dns_cache_table.next_addrs("foo")
         assert addrs == ["127.0.0.1"]
 
+    def test_max_size_eviction(self) -> None:
+        table = _DNSCacheTable(max_size=2)
+
+        table.add(self.host1, [self.result1])
+        table.add(self.host2, [self.result2])
+
+        host3 = ("example.com", 80)
+        result3: ResolveResult = {
+            **self.result1,
+            "hostname": "example.com",
+            "host": "1.2.3.4",
+        }
+        table.add(host3, [result3])
+
+        assert len(table._addrs_rr) == 2
+        assert self.host1 not in table._addrs_rr
+        assert host3 in table._addrs_rr
+
+    def test_lru_eviction(self) -> None:
+        table = _DNSCacheTable(max_size=2)
+
+        table.add(self.host1, [self.result1])
+        table.add(self.host2, [self.result2])
+
+        table.next_addrs(self.host1)
+
+        host3 = ("example.com", 80)
+        result3: ResolveResult = {
+            **self.result1,
+            "hostname": "example.com",
+            "host": "1.2.3.4",
+        }
+        table.add(host3, [result3])
+
+        assert self.host1 in table._addrs_rr
+        assert self.host2 not in table._addrs_rr
+
+    def test_lru_eviction_add(self) -> None:
+        table = _DNSCacheTable(max_size=2)
+
+        table.add(self.host1, [self.result1])
+        table.add(self.host2, [self.result2])
+
+        # Re-add, thus making host1 the most recently used.
+        table.add(self.host1, [self.result1])
+
+        host3 = ("example.com", 80)
+        result3: ResolveResult = {
+            **self.result1,
+            "hostname": "example.com",
+            "host": "1.2.3.4",
+        }
+        table.add(host3, [result3])
+
+        assert self.host1 in table._addrs_rr
+        assert self.host2 not in table._addrs_rr
+
 
 async def test_connector_cache_trace_race():
     class DummyTracer: