Merge pull request #3323 from lbryio/dht_leak

lbryio · Jun 7, 2021 · d861b08 · d861b08
2 parents b9142bb + fb438dc
commit d861b08
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 16 deletions.
diff --git a/lbry/dht/peer.py b/lbry/dht/peer.py
@@ -4,7 +4,7 @@
 from binascii import hexlify
 from dataclasses import dataclass, field
 from functools import lru_cache
-from lbry.utils import is_valid_public_ipv4 as _is_valid_public_ipv4
+from lbry.utils import is_valid_public_ipv4 as _is_valid_public_ipv4, LRUCache
 from lbry.dht import constants
 from lbry.dht.serialization.datagram import make_compact_address, make_compact_ip, decode_compact_address
 
@@ -31,12 +31,12 @@ def __init__(self, loop: asyncio.AbstractEventLoop):
         self._rpc_failures: typing.Dict[
             typing.Tuple[str, int], typing.Tuple[typing.Optional[float], typing.Optional[float]]
         ] = {}
-        self._last_replied: typing.Dict[typing.Tuple[str, int], float] = {}
-        self._last_sent: typing.Dict[typing.Tuple[str, int], float] = {}
-        self._last_requested: typing.Dict[typing.Tuple[str, int], float] = {}
-        self._node_id_mapping: typing.Dict[typing.Tuple[str, int], bytes] = {}
-        self._node_id_reverse_mapping: typing.Dict[bytes, typing.Tuple[str, int]] = {}
-        self._node_tokens: typing.Dict[bytes, (float, bytes)] = {}
+        self._last_replied: typing.Dict[typing.Tuple[str, int], float] = LRUCache(2048)
+        self._last_sent: typing.Dict[typing.Tuple[str, int], float] = LRUCache(2048)
+        self._last_requested: typing.Dict[typing.Tuple[str, int], float] = LRUCache(2048)
+        self._node_id_mapping: typing.Dict[typing.Tuple[str, int], bytes] = LRUCache(2048)
+        self._node_id_reverse_mapping: typing.Dict[bytes, typing.Tuple[str, int]] = LRUCache(2048)
+        self._node_tokens: typing.Dict[bytes, (float, bytes)] = LRUCache(2048)
 
     def reset(self):
         for statistic in (self._rpc_failures, self._last_replied, self._last_sent, self._last_requested):

diff --git a/lbry/utils.py b/lbry/utils.py
@@ -21,7 +21,6 @@
 import certifi
 import aiohttp
 from prometheus_client import Counter
-from prometheus_client.registry import REGISTRY
 from lbry.schema.claim import Claim
 
 
@@ -276,12 +275,6 @@ def __delitem__(self, key):
 
     def __del__(self):
         self.clear()
-        if self._track_metrics:  # needed for tests
-            try:
-                REGISTRY.unregister(self.hits)
-                REGISTRY.unregister(self.misses)
-            except AttributeError:
-                pass
 
 
 class LRUCache:
@@ -310,11 +303,14 @@ def set(self, key, value):
                 self.cache.popitem(last=False)
         self.cache[key] = value
 
+    def items(self):
+        return self.cache.items()
+
     def clear(self):
         self.cache.clear()
 
-    def pop(self, key):
-        return self.cache.pop(key)
+    def pop(self, key, default=None):
+        return self.cache.pop(key, default)
 
     def __setitem__(self, key, value):
         return self.set(key, value)