[WIP]

Author: thewhaleking

async_substrate_interface/async_substrate.py

@@ -21,7 +21,6 @@
     TYPE_CHECKING,
 )
 
-import asyncstdlib as a
 from bittensor_wallet.keypair import Keypair
 from bittensor_wallet.utils import SS58_FORMAT
 from bt_decode import MetadataV15, PortableRegistry, decode as decode_by_type_string
@@ -49,6 +48,7 @@
     Preprocessed,
 )
 from async_substrate_interface.utils import hex_to_bytes, json, get_next_id
+from async_substrate_interface.utils.cache import async_sql_lru_cache
 from async_substrate_interface.utils.decoding import (
     _determine_if_old_runtime_call,
     _bt_decode_to_dict_or_list,
@@ -1659,7 +1659,7 @@ def convert_event_data(data):
                 events.append(convert_event_data(item))
         return events
 
-    @a.lru_cache(maxsize=512)  # large cache with small items
+    @async_sql_lru_cache(max_size=512)
     async def get_parent_block_hash(self, block_hash):
         block_header = await self.rpc_request("chain_getHeader", [block_hash])
 
@@ -1672,15 +1672,15 @@ async def get_parent_block_hash(self, block_hash):
             return block_hash
         return parent_block_hash
 
-    @a.lru_cache(maxsize=16)  # small cache with large items
+    @async_sql_lru_cache(max_size=16)
     async def get_block_runtime_info(self, block_hash: str) -> dict:
         """
         Retrieve the runtime info of given block_hash
         """
         response = await self.rpc_request("state_getRuntimeVersion", [block_hash])
         return response.get("result")
 
-    @a.lru_cache(maxsize=512)  # large cache with small items
+    @async_sql_lru_cache(max_size=512)
     async def get_block_runtime_version_for(self, block_hash: str):
         """
         Retrieve the runtime version of the parent of a given block_hash
@@ -1914,7 +1914,7 @@ async def _make_rpc_request(
 
         return request_manager.get_results()
 
-    @a.lru_cache(maxsize=512)  # RPC methods are unlikely to change often
+    @async_sql_lru_cache(max_size=512)
     async def supports_rpc_method(self, name: str) -> bool:
         """
         Check if substrate RPC supports given method
@@ -1985,7 +1985,7 @@ async def rpc_request(
         else:
             raise SubstrateRequestException(result[payload_id][0])
 
-    @a.lru_cache(maxsize=512)  # block_id->block_hash does not change
+    @async_sql_lru_cache(max_size=512)
     async def get_block_hash(self, block_id: int) -> str:
         return (await self.rpc_request("chain_getBlockHash", [block_id]))["result"]
 

async_substrate_interface/sync_substrate.py

@@ -1,6 +1,5 @@
 import logging
 import random
-from functools import lru_cache
 from hashlib import blake2b
 from typing import Optional, Union, Callable, Any
 
@@ -31,6 +30,7 @@
     ScaleObj,
 )
 from async_substrate_interface.utils import hex_to_bytes, json, get_next_id
+from async_substrate_interface.utils.cache import sql_lru_cache
 from async_substrate_interface.utils.decoding import (
     _determine_if_old_runtime_call,
     _bt_decode_to_dict_or_list,
@@ -1406,7 +1406,7 @@ def convert_event_data(data):
                 events.append(convert_event_data(item))
         return events
 
-    @lru_cache(maxsize=512)  # large cache with small items
+    @sql_lru_cache(max_size=512)
     def get_parent_block_hash(self, block_hash):
         block_header = self.rpc_request("chain_getHeader", [block_hash])
 
@@ -1419,15 +1419,15 @@ def get_parent_block_hash(self, block_hash):
             return block_hash
         return parent_block_hash
 
-    @lru_cache(maxsize=16)  # small cache with large items
+    @sql_lru_cache(max_size=16)
     def get_block_runtime_info(self, block_hash: str) -> dict:
         """
         Retrieve the runtime info of given block_hash
         """
         response = self.rpc_request("state_getRuntimeVersion", [block_hash])
         return response.get("result")
 
-    @lru_cache(maxsize=512)  # large cache with small items
+    @sql_lru_cache(max_size=512)
     def get_block_runtime_version_for(self, block_hash: str):
         """
         Retrieve the runtime version of the parent of a given block_hash
@@ -1656,7 +1656,7 @@ def _make_rpc_request(
         return request_manager.get_results()
 
     # TODO change this logic
-    @lru_cache(maxsize=512)  # RPC methods are unlikely to change often
+    @sql_lru_cache(max_size=512)
     def supports_rpc_method(self, name: str) -> bool:
         """
         Check if substrate RPC supports given method
@@ -1727,7 +1727,7 @@ def rpc_request(
         else:
             raise SubstrateRequestException(result[payload_id][0])
 
-    @lru_cache(maxsize=512)  # block_id->block_hash does not change
+    @sql_lru_cache(max_size=512)
     def get_block_hash(self, block_id: int) -> str:
         return self.rpc_request("chain_getBlockHash", [block_id])["result"]
 

async_substrate_interface/utils/cache.py

@@ -0,0 +1,91 @@
+import functools
+import pickle
+import sqlite3
+import asyncstdlib as a
+
+
+def _get_table_name(func):
+    """Convert "ClassName.method_name" to "ClassName_method_name"""
+    return func.__qualname__.replace(".", "_")
+
+
+def _create_table(conn, table_name):
+    c = conn.cursor()
+    c.execute(
+        f"CREATE TABLE IF NOT EXISTS {table_name} (key BLOB PRIMARY KEY, value BLOB, chain TEXT)"
+    )
+    conn.commit()
+
+
+def _retrieve_from_cache(c, table_name, key, chain):
+    try:
+        c.execute(
+            f"SELECT value FROM {table_name} WHERE key=? AND chain=?", (key, chain)
+        )
+        result = c.fetchone()
+        if result is not None:
+            return pickle.loads(result[0])
+    except (pickle.PickleError, sqlite3.Error) as e:
+        print(f"Cache error: {str(e)}")
+        pass
+
+
+def _insert_into_cache(c, conn, table_name, key, result, chain):
+    try:
+        c.execute(
+            f"INSERT OR REPLACE INTO {table_name} VALUES (?,?,?)",
+            (key, pickle.dumps(result), chain),
+        )
+        conn.commit()
+    except (pickle.PickleError, sqlite3.Error) as e:
+        print(f"Cache error: {str(e)}")
+        pass
+
+
+def sql_lru_cache(func, max_size=None):
+    conn = sqlite3.connect("/tmp/cache.db")
+
+    table_name = _get_table_name(func)
+    _create_table(conn, table_name)
+
+    @functools.lru_cache(maxsize=max_size)
+    def inner(self, *args, **kwargs):
+        c = conn.cursor()
+        key = pickle.dumps((args, kwargs))
+        chain = self._chain
+
+        result = _retrieve_from_cache(c, table_name, key, chain)
+        if result is not None:
+            return result
+
+        # If not in DB, call func and store in DB
+        result = func(self, *args, **kwargs)
+        _insert_into_cache(c, conn, table_name, key, result, chain)
+
+        return result
+
+    return inner
+
+
+def async_sql_lru_cache(func, max_size=None):
+    conn = sqlite3.connect("/tmp/cache.db")
+    table_name = _get_table_name(func)
+    _create_table(conn, table_name)
+
+    @a.lru_cache(maxsize=max_size)
+    async def inner(self, *args, **kwargs):
+        c = conn.cursor()
+        key = pickle.dumps((args, kwargs))
+        chain = self._chain
+
+        result = _retrieve_from_cache(c, table_name, key, chain)
+        if result is not None:
+            return result
+
+        # If not in DB, call func and store in DB
+        result = await func(self, *args, **kwargs)
+        _insert_into_cache(c, conn, table_name, key, result, chain)
+
+        return result
+
+    return inner