chore: some storage adjustments (leanEthereum#427)

Author: tcoratger

src/lean_spec/subspecs/node/node.py

@@ -209,14 +209,20 @@ def from_genesis(cls, config: NodeConfig) -> Node:
             store = state.to_forkchoice_store(block, validator_id)
 
             # Persist genesis to database if available.
+            #
+            # Atomic write ensures a crash during genesis persistence
+            # does not leave the database in a partial state.
             if database is not None:
                 block_root = hash_tree_root(block)
-                database.put_block(block, block_root)
-                database.put_state(state, block_root)
-                database.put_head_root(block_root)
-                database.put_justified_checkpoint(store.latest_justified)
-                database.put_finalized_checkpoint(store.latest_finalized)
-                database.put_block_root_by_slot(block.slot, block_root)
+                with database.batch_write():
+                    database.put_block(block, block_root)
+                    database.put_state(state, block_root)
+                    database.put_head_root(block_root)
+                    database.put_justified_checkpoint(store.latest_justified)
+                    database.put_finalized_checkpoint(store.latest_finalized)
+                    database.put_block_root_by_slot(block.slot, block_root)
+                    database.put_block_root_by_state_root(hash_tree_root(state), block_root)
+                    database.put_genesis_time(config.genesis_time)
 
         # Create shared dependencies.
         clock = SlotClock(genesis_time=config.genesis_time, time_fn=config.time_fn)
@@ -353,6 +359,12 @@ def _try_load_from_database(
         if justified is None or finalized is None:
             return None
 
+        # Fall back to genesis time stored in the database.
+        #
+        # This enables self-contained restarts without external config.
+        if genesis_time is None:
+            genesis_time = database.get_genesis_time()
+
         # Compute store time from wall clock to avoid post-restart drift.
         #
         # Using only the head block's slot would set the store time to the

src/lean_spec/subspecs/storage/__init__.py

@@ -6,13 +6,14 @@
 """
 
 from .database import Database
-from .namespaces import BlockNamespace, CheckpointNamespace, StateNamespace
+from .exceptions import StorageCorruptionError, StorageError, StorageReadError, StorageWriteError
 from .sqlite import SQLiteDatabase
 
 __all__ = [
     "Database",
     "SQLiteDatabase",
-    "BlockNamespace",
-    "StateNamespace",
-    "CheckpointNamespace",
+    "StorageCorruptionError",
+    "StorageError",
+    "StorageReadError",
+    "StorageWriteError",
 ]

src/lean_spec/subspecs/storage/database.py

@@ -7,14 +7,16 @@
 
 from __future__ import annotations
 
+from collections.abc import Generator
+from contextlib import contextmanager
 from typing import TYPE_CHECKING, Protocol
 
 if TYPE_CHECKING:
     from lean_spec.subspecs.containers import Block, Checkpoint, State
     from lean_spec.subspecs.containers.attestation import AttestationData
     from lean_spec.subspecs.containers.slot import Slot
     from lean_spec.subspecs.containers.validator import ValidatorIndex
-    from lean_spec.types import Bytes32
+    from lean_spec.types import Bytes32, Uint64
 
 
 class Database(Protocol):
@@ -26,10 +28,11 @@ class Database(Protocol):
 
     Storage Organization
     --------------------
-    - Blocks: Indexed by root hash
-    - States: Indexed by root hash
+    - Blocks: Indexed by block root hash
+    - States: Indexed by associated block root hash (not state root)
     - Checkpoints: Justified and finalized tracking
     - Attestations: Latest attestation per validator
+    - State root index: Maps state roots to block roots
     """
 
     # Block Operations
@@ -72,10 +75,10 @@ def has_block(self, root: Bytes32) -> bool:
 
     def get_state(self, root: Bytes32) -> State | None:
         """
-        Retrieve a state by its root hash.
+        Retrieve a state by its associated block root.
 
         Args:
-            root: SSZ hash tree root of the state.
+            root: Block root hash associated with this state.
 
         Returns:
             State if found, None otherwise.
@@ -84,11 +87,11 @@ def get_state(self, root: Bytes32) -> State | None:
 
     def put_state(self, state: State, root: Bytes32) -> None:
         """
-        Store a state with its root hash.
+        Store a state indexed by its associated block root.
 
         Args:
             state: State to store.
-            root: Pre-computed root hash (avoids recomputation).
+            root: Block root hash associated with this state.
         """
         ...
 
@@ -97,7 +100,7 @@ def has_state(self, root: Bytes32) -> bool:
         Check if a state exists in storage.
 
         Args:
-            root: SSZ hash tree root of the state.
+            root: Block root hash associated with the state.
 
         Returns:
             True if state exists.
@@ -223,6 +226,93 @@ def put_block_root_by_slot(self, slot: Slot, root: Bytes32) -> None:
         """
         ...
 
+    # State Root Index Operations
+
+    def get_block_root_by_state_root(self, state_root: Bytes32) -> Bytes32 | None:
+        """
+        Look up the block root associated with a state root.
+
+        Needed for checkpoint sync and API endpoints that query by state root.
+
+        Args:
+            state_root: SSZ hash tree root of the state.
+
+        Returns:
+            Associated block root, or None if not indexed.
+        """
+        ...
+
+    def put_block_root_by_state_root(self, state_root: Bytes32, block_root: Bytes32) -> None:
+        """
+        Index a block root by the state root it produced.
+
+        Args:
+            state_root: SSZ hash tree root of the post-state.
+            block_root: Root of the block that produced this state.
+        """
+        ...
+
+    # Genesis Time
+
+    def get_genesis_time(self) -> Uint64 | None:
+        """
+        Retrieve the stored genesis time.
+
+        Enables self-contained restarts without external genesis config.
+
+        Returns:
+            Genesis time as Unix timestamp, or None if not set.
+        """
+        ...
+
+    def put_genesis_time(self, genesis_time: Uint64) -> None:
+        """
+        Store genesis time for future restarts.
+
+        Args:
+            genesis_time: Unix timestamp of genesis (slot 0).
+        """
+        ...
+
+    # Transaction Control
+
+    def commit(self) -> None:
+        """
+        Commit pending writes to durable storage.
+
+        All writes via put_* methods are buffered until commit() or batch_write().
+        Callers must explicitly commit after writes.
+        """
+        ...
+
+    @contextmanager
+    def batch_write(self) -> Generator[None]:
+        """
+        Context manager for atomic multi-write operations.
+
+        All writes within the block are committed atomically on exit.
+        Rolls back on exception to prevent partial writes.
+        """
+        ...
+
+    # Pruning
+
+    def prune_before_slot(self, slot: Slot, keep_roots: frozenset[Bytes32]) -> int:
+        """
+        Remove blocks and states with slots strictly before the given slot.
+
+        Preserves entries whose roots are in keep_roots (e.g., the finalized block).
+        Cleans up associated slot index entries.
+
+        Args:
+            slot: Prune entries with slots strictly below this value.
+            keep_roots: Roots to preserve regardless of slot.
+
+        Returns:
+            Total number of entries pruned across all tables.
+        """
+        ...
+
     # Lifecycle
 
     def close(self) -> None:

src/lean_spec/subspecs/storage/exceptions.py

@@ -0,0 +1,23 @@
+"""
+Storage exception hierarchy.
+
+Wraps low-level database and serialization errors with storage-specific context.
+This allows callers to handle storage failures uniformly without knowing the
+underlying backend (SQLite, etc.).
+"""
+
+
+class StorageError(Exception):
+    """Base exception for storage operations."""
+
+
+class StorageReadError(StorageError):
+    """Failed to read from storage."""
+
+
+class StorageWriteError(StorageError):
+    """Failed to write to storage."""
+
+
+class StorageCorruptionError(StorageError):
+    """Stored data failed deserialization."""

src/lean_spec/subspecs/storage/namespaces.py

@@ -86,6 +86,9 @@ class CheckpointNamespace:
     KEY_HEAD: str = "head"
     """Key for head block root."""
 
+    KEY_GENESIS_TIME: str = "genesis_time"
+    """Key for genesis time. Enables self-contained restarts without external config."""
+
     CREATE_TABLE: str = """
         CREATE TABLE IF NOT EXISTS checkpoints (
             key TEXT PRIMARY KEY,
@@ -137,9 +140,41 @@ class SlotIndexNamespace:
     """SQL to create slot index table."""
 
 
-# Singleton instances for convenient access
+@dataclass(frozen=True, slots=True)
+class StateRootIndexNamespace:
+    """
+    Namespace for state root to block root mapping.
+
+    Enables lookup of block root by state root.
+    Needed for checkpoint sync and API queries by state root.
+    """
+
+    TABLE_NAME: str = "state_root_index"
+    """Table name for state root index."""
+
+    CREATE_TABLE: str = """
+        CREATE TABLE IF NOT EXISTS state_root_index (
+            state_root BLOB PRIMARY KEY,
+            block_root BLOB NOT NULL
+        )
+    """
+    """SQL to create state root index table."""
+
+
 BLOCKS: Final = BlockNamespace()
+"""Block storage namespace."""
+
 STATES: Final = StateNamespace()
+"""State storage namespace."""
+
 CHECKPOINTS: Final = CheckpointNamespace()
+"""Checkpoint tracking namespace."""
+
 ATTESTATIONS: Final = AttestationNamespace()
+"""Attestation storage namespace."""
+
 SLOT_INDEX: Final = SlotIndexNamespace()
+"""Slot-to-root index namespace."""
+
+STATE_ROOT_INDEX: Final = StateRootIndexNamespace()
+"""State root to block root index namespace."""