Fixes potential cache issues.

Author: thewhaleking

async_substrate_interface/async_substrate.py

@@ -22,7 +22,6 @@
     TYPE_CHECKING,
 )
 
-import asyncstdlib as a
 from bt_decode import MetadataV15, PortableRegistry, decode as decode_by_type_string
 from scalecodec.base import ScaleBytes, ScaleType, RuntimeConfigurationObject
 from scalecodec.types import (
@@ -58,7 +57,7 @@
     get_next_id,
     rng as random,
 )
-from async_substrate_interface.utils.cache import async_sql_lru_cache
+from async_substrate_interface.utils.cache import async_sql_lru_cache, CachedFetcher
 from async_substrate_interface.utils.decoding import (
     _determine_if_old_runtime_call,
     _bt_decode_to_dict_or_list,
@@ -748,6 +747,12 @@ 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:
@@ -1869,9 +1874,8 @@ async def get_metadata(self, block_hash=None) -> MetadataV15:
 
         return runtime.metadata_v15
 
-    @a.lru_cache(maxsize=512)
     async def get_parent_block_hash(self, block_hash):
-        return await self._get_parent_block_hash(block_hash)
+        return await self._parent_hash_fetcher.execute(block_hash)
 
     async def _get_parent_block_hash(self, block_hash):
         block_header = await self.rpc_request("chain_getHeader", [block_hash])
@@ -1916,9 +1920,8 @@ async def get_storage_by_key(self, block_hash: str, storage_key: str) -> Any:
                 "Unknown error occurred during retrieval of events"
             )
 
-    @a.lru_cache(maxsize=16)
     async def get_block_runtime_info(self, block_hash: str) -> dict:
-        return await self._get_block_runtime_info(block_hash)
+        return await self._runtime_info_fetcher.execute(block_hash)
 
     get_block_runtime_version = get_block_runtime_info
 
@@ -1929,9 +1932,8 @@ async def _get_block_runtime_info(self, block_hash: str) -> dict:
         response = await self.rpc_request("state_getRuntimeVersion", [block_hash])
         return response.get("result")
 
-    @a.lru_cache(maxsize=512)
     async def get_block_runtime_version_for(self, block_hash: str):
-        return await self._get_block_runtime_version_for(block_hash)
+        return await self._runtime_version_for_fetcher.execute(block_hash)
 
     async def _get_block_runtime_version_for(self, block_hash: str):
         """
@@ -2240,9 +2242,8 @@ async def rpc_request(
         else:
             raise SubstrateRequestException(result[payload_id][0])
 
-    @a.lru_cache(maxsize=512)
     async def get_block_hash(self, block_id: int) -> str:
-        return await self._get_block_hash(block_id)
+        return await self._block_hash_fetcher.execute(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,10 +1,15 @@
+import asyncio
+from collections import OrderedDict
 import functools
 import os
 import pickle
 import sqlite3
 from pathlib import Path
+from typing import Callable, Any
+
 import asyncstdlib as a
 
+
 USE_CACHE = True if os.getenv("NO_CACHE") != "1" else False
 CACHE_LOCATION = (
     os.path.expanduser(
@@ -139,3 +144,54 @@ async def inner(self, *args, **kwargs):
         return inner
 
     return decorator
+
+
+class LRUCache:
+    def __init__(self, max_size: int):
+        self.max_size = max_size
+        self.cache = OrderedDict()
+
+    def set(self, key, value):
+        if key in self.cache:
+            self.cache.move_to_end(key)
+        self.cache[key] = value
+        if len(self.cache) > self.max_size:
+            self.cache.popitem(last=False)
+
+    def get(self, key):
+        if key in self.cache:
+            # Mark as recently used
+            self.cache.move_to_end(key)
+            return self.cache[key]
+        return None
+
+
+class CachedFetcher:
+    def __init__(self, max_size: int, method: Callable):
+        self._inflight: dict[int, asyncio.Future] = {}
+        self._method = method
+        self._cache = LRUCache(max_size=max_size)
+
+    async def execute(self, single_arg: Any) -> str:
+        if item := self._cache.get(single_arg):
+            return item
+
+        if single_arg in self._inflight:
+            result = await self._inflight[single_arg]
+            return result
+
+        loop = asyncio.get_running_loop()
+        future = loop.create_future()
+        self._inflight[single_arg] = future
+
+        try:
+            result = await self._method(single_arg)
+            self._cache.set(single_arg, 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)