From 13e38d6fd82f12363975dc2273c7a7b8a01ff207 Mon Sep 17 00:00:00 2001
From: Jack Robison <jackrobison@lbry.io>
Date: Wed, 23 Dec 2020 17:45:05 -0500
Subject: [PATCH] use LRUCache instead of pylru in wallet server

---
 lbry/utils.py                 | 6 +++---
 lbry/wallet/server/daemon.py  | 7 +++----
 lbry/wallet/server/leveldb.py | 5 ++---
 lbry/wallet/server/session.py | 7 +++----
 4 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/lbry/utils.py b/lbry/utils.py
index d449076f3..a2fd2cc86 100644
--- a/lbry/utils.py
+++ b/lbry/utils.py
@@ -215,7 +215,7 @@ class LRUCache:
         'misses'
     ]
 
-    def __init__(self, capacity: int, metric_name: typing.Optional[str] = None):
+    def __init__(self, capacity: int, metric_name: typing.Optional[str] = None, namespace: str = "daemon_cache"):
         self.capacity = capacity
         self.cache = collections.OrderedDict()
         if metric_name is None:
@@ -225,10 +225,10 @@ class LRUCache:
             self._track_metrics = True
             try:
                 self.hits = Counter(
-                    f"{metric_name}_cache_hit_count", "Number of cache hits", namespace="daemon_cache"
+                    f"{metric_name}_cache_hit_count", "Number of cache hits", namespace=namespace
                 )
                 self.misses = Counter(
-                    f"{metric_name}_cache_miss_count", "Number of cache misses", namespace="daemon_cache"
+                    f"{metric_name}_cache_miss_count", "Number of cache misses", namespace=namespace
                 )
             except ValueError as err:
                 log.warning("failed to set up prometheus %s_cache_miss_count metric: %s", metric_name, err)
diff --git a/lbry/wallet/server/daemon.py b/lbry/wallet/server/daemon.py
index 44a366b6a..654b4337e 100644
--- a/lbry/wallet/server/daemon.py
+++ b/lbry/wallet/server/daemon.py
@@ -3,11 +3,10 @@ import itertools
 import json
 import time
 from functools import wraps
-from pylru import lrucache
 
 import aiohttp
 from prometheus_client import Gauge, Histogram
-
+from lbry.utils import LRUCache
 from lbry.wallet.rpc.jsonrpc import RPCError
 from lbry.wallet.server.util import hex_to_bytes, class_logger
 from lbry.wallet.rpc import JSONRPC
@@ -55,8 +54,8 @@ class Daemon:
         self._height = None
         self.available_rpcs = {}
         self.connector = aiohttp.TCPConnector()
-        self._block_hash_cache = lrucache(100000)
-        self._block_cache = lrucache(10000)
+        self._block_hash_cache = LRUCache(100000)
+        self._block_cache = LRUCache(10000)
 
     async def close(self):
         if self.connector:
diff --git a/lbry/wallet/server/leveldb.py b/lbry/wallet/server/leveldb.py
index 50ac57a96..542f7d61f 100644
--- a/lbry/wallet/server/leveldb.py
+++ b/lbry/wallet/server/leveldb.py
@@ -15,7 +15,6 @@ import ast
 import os
 import time
 import zlib
-import pylru
 import typing
 from typing import Optional, List, Tuple, Iterable
 from asyncio import sleep
@@ -25,7 +24,7 @@ from glob import glob
 from struct import pack, unpack
 from concurrent.futures.thread import ThreadPoolExecutor
 import attr
-
+from lbry.utils import LRUCache
 from lbry.wallet.server import util
 from lbry.wallet.server.hash import hash_to_hex_str, HASHX_LEN
 from lbry.wallet.server.merkle import Merkle, MerkleCache
@@ -94,7 +93,7 @@ class LevelDB:
         self.headers_db = None
         self.tx_db = None
 
-        self._tx_and_merkle_cache = pylru.lrucache(100000)
+        self._tx_and_merkle_cache = LRUCache(100000, metric_name='tx_and_merkle', namespace="wallet_server")
         self.total_transactions = None
 
     async def _read_tx_counts(self):
diff --git a/lbry/wallet/server/session.py b/lbry/wallet/server/session.py
index 6f9808497..1436492eb 100644
--- a/lbry/wallet/server/session.py
+++ b/lbry/wallet/server/session.py
@@ -4,7 +4,6 @@ import math
 import time
 import json
 import zlib
-import pylru
 import base64
 import codecs
 import typing
@@ -18,11 +17,11 @@ from collections import defaultdict
 from functools import partial
 
 from binascii import hexlify, unhexlify
-from pylru import lrucache
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 from prometheus_client import Counter, Info, Histogram, Gauge
 
 import lbry
+from lbry.utils import LRUCache
 from lbry.build_info import BUILD, COMMIT_HASH, DOCKER_TAG
 from lbry.wallet.server.block_processor import LBRYBlockProcessor
 from lbry.wallet.server.db.writer import LBRYLevelDB
@@ -811,8 +810,8 @@ class LBRYSessionManager(SessionManager):
         if self.env.websocket_host is not None and self.env.websocket_port is not None:
             self.websocket = AdminWebSocket(self)
         self.search_cache = self.bp.search_cache
-        self.search_cache['search'] = lrucache(10000)
-        self.search_cache['resolve'] = lrucache(10000)
+        self.search_cache['search'] = LRUCache(10000, metric_name='search', namespace=NAMESPACE)
+        self.search_cache['resolve'] = LRUCache(10000, metric_name='resolve', namespace=NAMESPACE)
 
     async def process_metrics(self):
         while self.running: