WIP

thewhaleking · thewhaleking · commit ea08a69a1ba3 · 2025-10-27T23:09:25.000+02:00
diff --git a/async_substrate_interface/async_substrate.py b/async_substrate_interface/async_substrate.py
@@ -1091,7 +1091,10 @@ async def __aenter__(self):
             await self.initialize()
         return self
 
-    async def initialize(self):
+    async def initialize(self) -> None:
+        await self._initialize()
+
+    async def _initialize(self) -> None:
         """
         Initialize the connection to the chain.
         """
@@ -1116,7 +1119,7 @@ async def initialize(self):
         self._initializing = False
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.ws.shutdown()
+        await self.close()
 
     @property
     def metadata(self):
@@ -4260,11 +4263,16 @@ class DiskCachedAsyncSubstrateInterface(AsyncSubstrateInterface):
     Experimental new class that uses disk-caching in addition to memory-caching for the cached methods
     """
 
+    async def initialize(self) -> None:
+        await self.runtime_cache.load_from_disk(self.url)
+        await self._initialize()
+
     async def close(self):
         """
         Closes the substrate connection, and the websocket connection.
         """
         try:
+            await self.runtime_cache.dump_to_disk(self.url)
             await self.ws.shutdown()
         except AttributeError:
             pass
diff --git a/async_substrate_interface/types.py b/async_substrate_interface/types.py
@@ -16,6 +16,7 @@
 
 from .const import SS58_FORMAT
 from .utils import json
+from .utils.cache import AsyncSqliteDB
 
 logger = logging.getLogger("async_substrate_interface")
 
@@ -101,6 +102,21 @@ def retrieve(
                 return runtime
         return None
 
+    async def load_from_disk(self, chain_endpoint: str):
+        db = AsyncSqliteDB(chain_endpoint=chain_endpoint)
+        block_mapping, block_hash_mapping, runtime_version_mapping = await db.load_runtime_cache(chain_endpoint)
+        if not any([block_mapping, block_hash_mapping, runtime_version_mapping]):
+            logger.debug("No runtime mappings in disk cache")
+        else:
+            logger.debug("Found runtime mappings in disk cache")
+        self.blocks = {x: Runtime.deserialize(y) for x, y in block_mapping.items()}
+        self.block_hashes = {x: Runtime.deserialize(y) for x, y in block_hash_mapping.items()}
+        self.versions = {x: Runtime.deserialize(y) for x, y in runtime_version_mapping.items()}
+
+    async def dump_to_disk(self, chain_endpoint: str):
+        db = AsyncSqliteDB(chain_endpoint=chain_endpoint)
+        await db.dump_runtime_cache(chain_endpoint, self.blocks, self.block_hashes, self.versions)
+
 
 class Runtime:
     """
@@ -149,6 +165,33 @@ def __init__(
         if registry is not None:
             self.load_registry_type_map()
 
+    def serialize(self):
+        return {
+            "chain": self.chain,
+            "type_registry": self.type_registry,
+            "metadata": self.metadata,
+            "metadata_v15": self.metadata_v15.to_json() if self.metadata_v15 is not None else None,
+            "runtime_info": self.runtime_info,
+            "registry": None,  # gets loaded from metadata_v15
+            "ss58_format": self.ss58_format,
+            "runtime_config": self.runtime_config,
+        }
+
+    @classmethod
+    def deserialize(cls, serialized: dict) -> "Runtime":
+        mdv15 = MetadataV15
+        registry = PortableRegistry.from_metadata_v15(mdv15) if (mdv15 := serialized["metadata_v15"]) else None
+        return cls(
+            chain=serialized["chain"],
+            metadata=serialized["metadata"],
+            type_registry=serialized["type_registry"],
+            runtime_config=serialized["runtime_config"],
+            metadata_v15=mdv15 if mdv15 is not None else None,
+            registry=registry,
+            ss58_format=serialized["ss58_format"],
+            runtime_info=serialized["runtime_info"]
+        )
+
     def load_runtime(self):
         """
         Initial loading of the runtime's type registry information.
diff --git a/async_substrate_interface/utils/cache.py b/async_substrate_interface/utils/cache.py
@@ -4,12 +4,12 @@
 import functools
 import logging
 import os
-import pickle
 import sqlite3
 from pathlib import Path
 from typing import Callable, Any, Awaitable, Hashable, Optional
 
 import aiosqlite
+import dill as pickle
 
 
 USE_CACHE = True if os.getenv("NO_CACHE") != "1" else False
@@ -38,13 +38,11 @@ def __new__(cls, chain_endpoint: str):
             cls._instances[chain_endpoint] = instance
             return instance
 
-    async def __call__(self, chain, other_self, func, args, kwargs) -> Optional[Any]:
+    async def _create_if_not_exists(self, chain: str, table_name: str):
         async with self._lock:
             if not self._db:
                 _ensure_dir()
                 self._db = await aiosqlite.connect(CACHE_LOCATION)
-        table_name = _get_table_name(func)
-        key = None
         if not (local_chain := _check_if_local(chain)) or not USE_CACHE:
             await self._db.execute(
                 f"""
@@ -72,19 +70,24 @@ async def __call__(self, chain, other_self, func, args, kwargs) -> Optional[Any]
                 """
             )
             await self._db.commit()
-            key = pickle.dumps((args, kwargs or None))
-            try:
-                cursor: aiosqlite.Cursor = await self._db.execute(
-                    f"SELECT value FROM {table_name} WHERE key=? AND chain=?",
-                    (key, chain),
-                )
-                result = await cursor.fetchone()
-                await cursor.close()
-                if result is not None:
-                    return pickle.loads(result[0])
-            except (pickle.PickleError, sqlite3.Error) as e:
-                logger.exception("Cache error", exc_info=e)
-                pass
+        return local_chain
+
+    async def __call__(self, chain, other_self, func, args, kwargs) -> Optional[Any]:
+        table_name = _get_table_name(func)
+        local_chain = await self._create_if_not_exists(chain, table_name)
+        key = pickle.dumps((args, kwargs or None))
+        try:
+            cursor: aiosqlite.Cursor = await self._db.execute(
+                f"SELECT value FROM {table_name} WHERE key=? AND chain=?",
+                (key, chain),
+            )
+            result = await cursor.fetchone()
+            await cursor.close()
+            if result is not None:
+                return pickle.loads(result[0])
+        except (pickle.PickleError, sqlite3.Error) as e:
+            logger.exception("Cache error", exc_info=e)
+            pass
         result = await func(other_self, *args, **kwargs)
         if not local_chain or not USE_CACHE:
             # TODO use a task here
@@ -95,6 +98,59 @@ async def __call__(self, chain, other_self, func, args, kwargs) -> Optional[Any]
             await self._db.commit()
         return result
 
+    async def load_runtime_cache(self, chain: str) -> tuple[dict, dict, dict]:
+        block_mapping = {}
+        block_hash_mapping = {}
+        version_mapping = {}
+        tables = {
+            "rt_cache_block": block_mapping,
+            "rt_cache_block_hash": block_hash_mapping,
+            "rt_cache_version": version_mapping
+        }
+        for table in tables.keys():
+            local_chain = await self._create_if_not_exists(chain, table)
+            if local_chain:
+                return {}, {}, {}
+        for table_name, mapping in tables.items():
+            try:
+                cursor: aiosqlite.Cursor = await self._db.execute(
+                    f"SELECT key, value FROM {table_name} WHERE chain=?",
+                    (chain,),
+                )
+                results = await cursor.fetchall()
+                await cursor.close()
+                if results is None:
+                    continue
+                for row in results:
+                    key, value = row
+                    runtime = pickle.loads(value)
+                    mapping[key] = runtime
+            except (pickle.PickleError, sqlite3.Error) as e:
+                logger.exception("Cache error", exc_info=e)
+                return {}, {}, {}
+        return block_mapping, block_hash_mapping, version_mapping
+
+    async def dump_runtime_cache(self, chain: str, block_mapping: dict, block_hash_mapping: dict, version_mapping: dict) -> None:
+        async with self._lock:
+            if not self._db:
+                _ensure_dir()
+                self._db = await aiosqlite.connect(CACHE_LOCATION)
+        tables = {
+            "rt_cache_block": block_mapping,
+            "rt_cache_block_hash": block_hash_mapping,
+            "rt_cache_version": version_mapping
+        }
+        for table, mapping in tables.items():
+            local_chain = await self._create_if_not_exists(chain, table)
+            if local_chain:
+                return None
+            await self._db.executemany(
+                f"INSERT OR REPLACE INTO {table} (key, value, chain) VALUES (?,?,?)",
+                [(key, pickle.dumps(runtime.serialize()), chain) for key, runtime in mapping.items()],
+            )
+            await self._db.commit()
+        return None
+
 
 def _ensure_dir():
     path = Path(CACHE_LOCATION).parent
