Improves CachedFetcher

Author: thewhaleking

async_substrate_interface/async_substrate.py

@@ -58,7 +58,11 @@
     get_next_id,
     rng as random,
 )
-from async_substrate_interface.utils.cache import async_sql_lru_cache, CachedFetcher
+from async_substrate_interface.utils.cache import (
+    async_sql_lru_cache,
+    CachedFetcher,
+    cached_fetcher,
+)
 from async_substrate_interface.utils.decoding import (
     _determine_if_old_runtime_call,
     _bt_decode_to_dict_or_list,
@@ -794,12 +798,6 @@ def __init__(
         self.registry_type_map = {}
         self.type_id_to_name = {}
         self._mock = _mock
-        self._block_hash_fetcher = CachedFetcher(512, self._get_block_hash)
-        self._parent_hash_fetcher = CachedFetcher(512, self._get_parent_block_hash)
-        self._runtime_info_fetcher = CachedFetcher(16, self._get_block_runtime_info)
-        self._runtime_version_for_fetcher = CachedFetcher(
-            512, self._get_block_runtime_version_for
-        )
 
     async def __aenter__(self):
         if not self._mock:
@@ -1044,35 +1042,7 @@ async def init_runtime(
         if not runtime:
             self.last_block_hash = block_hash
 
-            runtime_block_hash = await self.get_parent_block_hash(block_hash)
-
-            runtime_info = await self.get_block_runtime_info(runtime_block_hash)
-
-            metadata, (metadata_v15, registry) = await asyncio.gather(
-                self.get_block_metadata(block_hash=runtime_block_hash, decode=True),
-                self._load_registry_at_block(block_hash=runtime_block_hash),
-            )
-            if metadata is None:
-                # does this ever happen?
-                raise SubstrateRequestException(
-                    f"No metadata for block '{runtime_block_hash}'"
-                )
-            logger.debug(
-                f"Retrieved metadata and metadata v15 for {runtime_version} from Substrate node"
-            )
-
-            runtime = Runtime(
-                chain=self.chain,
-                runtime_config=self.runtime_config,
-                metadata=metadata,
-                type_registry=self.type_registry,
-                metadata_v15=metadata_v15,
-                runtime_info=runtime_info,
-                registry=registry,
-            )
-            self.runtime_cache.add_item(
-                runtime_version=runtime_version, runtime=runtime
-            )
+            runtime = await self.get_runtime_for_version(runtime_version, block_hash)
 
         self.load_runtime(runtime)
 
@@ -1086,6 +1056,42 @@ async def init_runtime(
                 self.ss58_format = ss58_prefix_constant
         return runtime
 
+    @cached_fetcher(max_size=16, cache_key_index=0)
+    async def get_runtime_for_version(
+        self, runtime_version: int, block_hash: Optional[str] = None
+    ) -> Runtime:
+        return await self._get_runtime_for_version(runtime_version, block_hash)
+
+    async def _get_runtime_for_version(
+        self, runtime_version: int, block_hash: Optional[str] = None
+    ) -> Runtime:
+        runtime_block_hash = await self.get_parent_block_hash(block_hash)
+        runtime_info, metadata, (metadata_v15, registry) = await asyncio.gather(
+            self.get_block_runtime_info(runtime_block_hash),
+            self.get_block_metadata(block_hash=runtime_block_hash, decode=True),
+            self._load_registry_at_block(block_hash=runtime_block_hash),
+        )
+        if metadata is None:
+            # does this ever happen?
+            raise SubstrateRequestException(
+                f"No metadata for block '{runtime_block_hash}'"
+            )
+        logger.debug(
+            f"Retrieved metadata and metadata v15 for {runtime_version} from Substrate node"
+        )
+
+        runtime = Runtime(
+            chain=self.chain,
+            runtime_config=self.runtime_config,
+            metadata=metadata,
+            type_registry=self.type_registry,
+            metadata_v15=metadata_v15,
+            runtime_info=runtime_info,
+            registry=registry,
+        )
+        self.runtime_cache.add_item(runtime_version=runtime_version, runtime=runtime)
+        return runtime
+
     async def create_storage_key(
         self,
         pallet: str,
@@ -1921,8 +1927,9 @@ async def get_metadata(self, block_hash=None) -> MetadataV15:
 
         return runtime.metadata_v15
 
+    @cached_fetcher(max_size=512)
     async def get_parent_block_hash(self, block_hash):
-        return await self._parent_hash_fetcher.execute(block_hash)
+        return await self._get_parent_block_hash(block_hash)
 
     async def _get_parent_block_hash(self, block_hash):
         block_header = await self.rpc_request("chain_getHeader", [block_hash])
@@ -1967,8 +1974,9 @@ async def get_storage_by_key(self, block_hash: str, storage_key: str) -> Any:
                 "Unknown error occurred during retrieval of events"
             )
 
+    @cached_fetcher(max_size=16)
     async def get_block_runtime_info(self, block_hash: str) -> dict:
-        return await self._runtime_info_fetcher.execute(block_hash)
+        return await self._get_block_runtime_info(block_hash)
 
     get_block_runtime_version = get_block_runtime_info
 
@@ -1979,8 +1987,9 @@ async def _get_block_runtime_info(self, block_hash: str) -> dict:
         response = await self.rpc_request("state_getRuntimeVersion", [block_hash])
         return response.get("result")
 
+    @cached_fetcher(max_size=512)
     async def get_block_runtime_version_for(self, block_hash: str):
-        return await self._runtime_version_for_fetcher.execute(block_hash)
+        return await self._get_block_runtime_version_for(block_hash)
 
     async def _get_block_runtime_version_for(self, block_hash: str):
         """
@@ -2296,8 +2305,9 @@ async def rpc_request(
         else:
             raise SubstrateRequestException(result[payload_id][0])
 
+    @cached_fetcher(max_size=512)
     async def get_block_hash(self, block_id: int) -> str:
-        return await self._block_hash_fetcher.execute(block_id)
+        return await self._get_block_hash(block_id)
 
     async def _get_block_hash(self, block_id: int) -> str:
         return (await self.rpc_request("chain_getBlockHash", [block_id]))["result"]

async_substrate_interface/utils/cache.py

@@ -1,13 +1,12 @@
 import asyncio
 from collections import OrderedDict
 import functools
+import logging
 import os
 import pickle
 import sqlite3
 from pathlib import Path
-from typing import Callable, Any
-
-import asyncstdlib as a
+from typing import Callable, Any, Awaitable, Hashable
 
 
 USE_CACHE = True if os.getenv("NO_CACHE") != "1" else False
@@ -19,6 +18,8 @@
     else ":memory:"
 )
 
+logger = logging.getLogger("async_substrate_interface")
+
 
 def _ensure_dir():
     path = Path(CACHE_LOCATION).parent
@@ -70,7 +71,7 @@ def _retrieve_from_cache(c, table_name, key, chain):
         if result is not None:
             return pickle.loads(result[0])
     except (pickle.PickleError, sqlite3.Error) as e:
-        print(f"Cache error: {str(e)}")
+        logger.exception("Cache error", exc_info=e)
         pass
 
 
@@ -82,7 +83,7 @@ def _insert_into_cache(c, conn, table_name, key, result, chain):
         )
         conn.commit()
     except (pickle.PickleError, sqlite3.Error) as e:
-        print(f"Cache error: {str(e)}")
+        logger.exception("Cache error", exc_info=e)
         pass
 
 
@@ -128,7 +129,7 @@ def inner(self, *args, **kwargs):
 
 def async_sql_lru_cache(maxsize=None):
     def decorator(func):
-        @a.lru_cache(maxsize=maxsize)
+        @cached_fetcher(max_size=maxsize)
         async def inner(self, *args, **kwargs):
             c, conn, table_name, key, result, chain, local_chain = (
                 _shared_inner_fn_logic(func, self, args, kwargs)
@@ -167,31 +168,65 @@ def get(self, key):
 
 
 class CachedFetcher:
-    def __init__(self, max_size: int, method: Callable):
-        self._inflight: dict[int, asyncio.Future] = {}
+    def __init__(
+        self,
+        max_size: int,
+        method: Callable[..., Awaitable[Any]],
+        cache_key_index: int = 0,
+    ):
+        self._inflight: dict[Hashable, asyncio.Future] = {}
         self._method = method
         self._cache = LRUCache(max_size=max_size)
+        self._cache_key_index = cache_key_index
 
-    async def execute(self, single_arg: Any) -> str:
-        if item := self._cache.get(single_arg):
+    async def __call__(self, *args: Any) -> Any:
+        key = args[self._cache_key_index]
+        if item := self._cache.get(key):
             return item
 
-        if single_arg in self._inflight:
-            result = await self._inflight[single_arg]
-            return result
+        if key in self._inflight:
+            return await self._inflight[key]
 
         loop = asyncio.get_running_loop()
         future = loop.create_future()
-        self._inflight[single_arg] = future
+        self._inflight[key] = future
 
         try:
-            result = await self._method(single_arg)
-            self._cache.set(single_arg, result)
+            result = await self._method(*args)
+            self._cache.set(key, result)
             future.set_result(result)
             return result
         except Exception as e:
-            # Propagate errors
             future.set_exception(e)
             raise
         finally:
-            self._inflight.pop(single_arg, None)
+            self._inflight.pop(key, None)
+
+
+class CachedFetcherMethod:
+    def __init__(self, method, max_size: int, cache_key_index: int):
+        self.method = method
+        self.max_size = max_size
+        self.cache_key_index = cache_key_index
+        self._instances = {}
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            return self
+
+        # Cache per-instance
+        if instance not in self._instances:
+            bound_method = self.method.__get__(instance, owner)
+            self._instances[instance] = CachedFetcher(
+                max_size=self.max_size,
+                method=bound_method,
+                cache_key_index=self.cache_key_index,
+            )
+        return self._instances[instance]
+
+
+def cached_fetcher(max_size: int, cache_key_index: int = 0):
+    def wrapper(method):
+        return CachedFetcherMethod(method, max_size, cache_key_index)
+
+    return wrapper

pyproject.toml

@@ -8,7 +8,6 @@ keywords = ["substrate", "development", "bittensor"]
 
 dependencies = [
   "wheel",
-  "asyncstdlib~=3.13.0",
   "bt-decode==v0.6.0",
   "scalecodec~=1.2.11",
   "websockets>=14.1",

tests/unit_tests/test_cache.py

@@ -13,18 +13,18 @@ async def test_cached_fetcher_fetches_and_caches():
     fetcher = CachedFetcher(max_size=2, method=mock_method)
 
     # First call should trigger the method
-    result1 = await fetcher.execute("key1")
+    result1 = await fetcher("key1")
     assert result1 == "result_key1"
     mock_method.assert_awaited_once_with("key1")
 
     # Second call with the same key should use the cache
-    result2 = await fetcher.execute("key1")
+    result2 = await fetcher("key1")
     assert result2 == "result_key1"
     # Ensure the method was NOT called again
     assert mock_method.await_count == 1
 
     # Third call with a new key triggers a method call
-    result3 = await fetcher.execute("key2")
+    result3 = await fetcher("key2")
     assert result3 == "result_key2"
     assert mock_method.await_count == 2
 
@@ -42,11 +42,11 @@ async def slow_method(x):
     fetcher = CachedFetcher(max_size=2, method=slow_method)
 
     # Start first request
-    task1 = asyncio.create_task(fetcher.execute("key1"))
+    task1 = asyncio.create_task(fetcher("key1"))
     await asyncio.sleep(0.1)  # Let the task start and be inflight
 
     # Second request for the same key while the first is in-flight
-    task2 = asyncio.create_task(fetcher.execute("key1"))
+    task2 = asyncio.create_task(fetcher("key1"))
     await asyncio.sleep(0.1)
 
     # Release the inflight request
@@ -65,7 +65,7 @@ async def error_method(x):
     fetcher = CachedFetcher(max_size=2, method=error_method)
 
     with pytest.raises(ValueError, match="Boom!"):
-        await fetcher.execute("key1")
+        await fetcher("key1")
 
 
 @pytest.mark.asyncio
@@ -75,12 +75,12 @@ async def test_cached_fetcher_eviction():
     fetcher = CachedFetcher(max_size=2, method=mock_method)
 
     # Fill cache
-    await fetcher.execute("key1")
-    await fetcher.execute("key2")
+    await fetcher("key1")
+    await fetcher("key2")
     assert list(fetcher._cache.cache.keys()) == list(fetcher._cache.cache.keys())
 
     # Insert a new key to trigger eviction
-    await fetcher.execute("key3")
+    await fetcher("key3")
     # key1 should be evicted
     assert "key1" not in fetcher._cache.cache
     assert "key2" in fetcher._cache.cache