refactor(storage): create StorageProtocol (#922)

Author: glevco

hathor/daa.py

@@ -27,10 +27,11 @@
 
 from hathor.conf.settings import HathorSettings
 from hathor.profiler import get_cpu_profiler
-from hathor.util import iwindows
+from hathor.util import iwindows, not_none
 
 if TYPE_CHECKING:
     from hathor.transaction import Block, Transaction
+    from hathor.transaction.storage.vertex_storage_protocol import VertexStorageProtocol
 
 logger = get_logger()
 cpu = get_cpu_profiler()
@@ -65,9 +66,9 @@ def calculate_block_difficulty(self, block: 'Block') -> float:
         if block.is_genesis:
             return self.MIN_BLOCK_WEIGHT
 
-        return self.calculate_next_weight(block.get_block_parent(), block.timestamp)
+        return self.calculate_next_weight(block.get_block_parent(), block.timestamp, not_none(block.storage))
 
-    def calculate_next_weight(self, parent_block: 'Block', timestamp: int) -> float:
+    def calculate_next_weight(self, parent_block: 'Block', timestamp: int, storage: 'VertexStorageProtocol') -> float:
         """ Calculate the next block weight, aka DAA/difficulty adjustment algorithm.
 
         The algorithm used is described in [RFC 22](https://gitlab.com/HathorNetwork/rfcs/merge_requests/22).
@@ -90,7 +91,7 @@ def calculate_next_weight(self, parent_block: 'Block', timestamp: int) -> float:
         blocks: list['Block'] = []
         while len(blocks) < N + 1:
             blocks.append(root)
-            root = root.get_block_parent()
+            root = storage.get_parent_block(root)
             assert root is not None
 
         # TODO: revise if this assertion can be safely removed

hathor/manager.py

@@ -51,6 +51,7 @@
 from hathor.profiler import get_cpu_profiler
 from hathor.pubsub import HathorEvents, PubSubManager
 from hathor.reactor import ReactorProtocol as Reactor
+from hathor.reward_lock import is_spent_reward_locked
 from hathor.stratum import StratumFactory
 from hathor.transaction import BaseTransaction, Block, MergeMinedBlock, Transaction, TxVersion, sum_weights
 from hathor.transaction.exceptions import TxValidationError
@@ -802,7 +803,7 @@ def _make_block_template(self, parent_block: Block, parent_txs: 'ParentTxs', cur
             parent_block_metadata.score,
             2 * self._settings.WEIGHT_TOL
         )
-        weight = max(self.daa.calculate_next_weight(parent_block, timestamp), min_significant_weight)
+        weight = max(self.daa.calculate_next_weight(parent_block, timestamp, self.tx_storage), min_significant_weight)
         height = parent_block.get_height() + 1
         parents = [parent_block.hash] + parent_txs.must_include
         parents_any = parent_txs.can_include
@@ -889,8 +890,7 @@ def push_tx(self, tx: Transaction, allow_non_standard_script: bool = False,
         if is_spending_voided_tx:
             raise SpendingVoidedError('Invalid transaction. At least one input is voided.')
 
-        is_spent_reward_locked = tx.is_spent_reward_locked()
-        if is_spent_reward_locked:
+        if is_spent_reward_locked(tx):
             raise RewardLockedError('Spent reward is locked.')
 
         # We are using here the method from lib because the property

hathor/reward_lock/__init__.py

@@ -0,0 +1,21 @@
+#  Copyright 2024 Hathor Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from hathor.reward_lock.reward_lock import get_spent_reward_locked_info, is_spent_reward_locked, iter_spent_rewards
+
+__all__ = [
+    'iter_spent_rewards',
+    'is_spent_reward_locked',
+    'get_spent_reward_locked_info',
+]

hathor/reward_lock/reward_lock.py

@@ -0,0 +1,69 @@
+#  Copyright 2024 Hathor Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from typing import TYPE_CHECKING, Iterator, Optional
+
+from hathor.conf.get_settings import get_global_settings
+from hathor.transaction import Block
+from hathor.util import not_none
+
+if TYPE_CHECKING:
+    from hathor.transaction.storage.vertex_storage_protocol import VertexStorageProtocol
+    from hathor.transaction.transaction import RewardLockedInfo, Transaction
+
+
+def iter_spent_rewards(tx: 'Transaction', storage: 'VertexStorageProtocol') -> Iterator[Block]:
+    """Iterate over all the rewards being spent, assumes tx has been verified."""
+    for input_tx in tx.inputs:
+        spent_tx = storage.get_vertex(input_tx.tx_id)
+        if spent_tx.is_block:
+            assert isinstance(spent_tx, Block)
+            yield spent_tx
+
+
+def is_spent_reward_locked(tx: 'Transaction') -> bool:
+    """ Check whether any spent reward is currently locked, considering only the block rewards spent by this tx
+    itself, and not the inherited `min_height`"""
+    return get_spent_reward_locked_info(tx, not_none(tx.storage)) is not None
+
+
+def get_spent_reward_locked_info(tx: 'Transaction', storage: 'VertexStorageProtocol') -> Optional['RewardLockedInfo']:
+    """Check if any input block reward is locked, returning the locked information if any, or None if they are all
+    unlocked."""
+    from hathor.transaction.transaction import RewardLockedInfo
+    for blk in iter_spent_rewards(tx, storage):
+        assert blk.hash is not None
+        needed_height = _spent_reward_needed_height(blk, storage)
+        if needed_height > 0:
+            return RewardLockedInfo(blk.hash, needed_height)
+    return None
+
+
+def _spent_reward_needed_height(block: Block, storage: 'VertexStorageProtocol') -> int:
+    """ Returns height still needed to unlock this `block` reward: 0 means it's unlocked."""
+    import math
+
+    # omitting timestamp to get the current best block, this will usually hit the cache instead of being slow
+    tips = storage.get_best_block_tips()
+    assert len(tips) > 0
+    best_height = math.inf
+    for tip in tips:
+        blk = storage.get_block(tip)
+        best_height = min(best_height, blk.get_height())
+    assert isinstance(best_height, int)
+    spent_height = block.get_height()
+    spend_blocks = best_height - spent_height
+    settings = get_global_settings()
+    needed_height = settings.REWARD_SPEND_MIN_BLOCKS - spend_blocks
+    return max(needed_height, 0)

hathor/transaction/storage/transaction_storage.py

@@ -45,6 +45,7 @@
 from hathor.transaction.storage.tx_allow_scope import TxAllowScope, tx_allow_context
 from hathor.transaction.transaction import Transaction
 from hathor.transaction.transaction_metadata import TransactionMetadata
+from hathor.types import VertexId
 from hathor.util import not_none
 
 cpu = get_cpu_profiler()
@@ -1137,6 +1138,17 @@ def _construct_genesis_tx2(self) -> Transaction:
         assert tx2.hash == self._settings.GENESIS_TX2_HASH
         return tx2
 
+    def get_parent_block(self, block: Block) -> Block:
+        return block.get_block_parent()
+
+    def get_vertex(self, vertex_id: VertexId) -> BaseTransaction:
+        return self.get_transaction(vertex_id)
+
+    def get_block(self, block_id: VertexId) -> Block:
+        block = self.get_vertex(block_id)
+        assert isinstance(block, Block)
+        return block
+
 
 class BaseTransactionStorage(TransactionStorage):
     indexes: Optional[IndexesManager]

hathor/transaction/storage/vertex_storage_protocol.py

@@ -0,0 +1,48 @@
+#  Copyright 2024 Hathor Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from abc import abstractmethod
+from typing import Protocol
+
+from hathor.transaction import BaseTransaction, Block
+from hathor.types import VertexId
+
+
+class VertexStorageProtocol(Protocol):
+    """
+    This Protocol currently represents a subset of TransactionStorage methods. Its main use case is for verification
+    methods that can receive a RocksDB storage or an ephemeral simple memory storage.
+
+    Therefore, objects returned by this protocol may or may not have an `object.storage` pointer set.
+    """
+
+    @abstractmethod
+    def get_vertex(self, vertex_id: VertexId) -> BaseTransaction:
+        """Return a vertex from the storage."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_block(self, block_id: VertexId) -> Block:
+        """Return a block from the storage."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_parent_block(self, block: Block) -> Block:
+        """Get the parent block of a block."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_best_block_tips(self) -> list[VertexId]:
+        """Return a list of blocks that are heads in a best chain."""
+        raise NotImplementedError

hathor/transaction/transaction.py

@@ -15,12 +15,13 @@
 import hashlib
 from itertools import chain
 from struct import pack
-from typing import TYPE_CHECKING, Any, Iterator, NamedTuple, Optional
+from typing import TYPE_CHECKING, Any, NamedTuple, Optional
 
 from hathor.checkpoint import Checkpoint
 from hathor.exception import InvalidNewTransaction
 from hathor.profiler import get_cpu_profiler
-from hathor.transaction import BaseTransaction, Block, TxInput, TxOutput, TxVersion
+from hathor.reward_lock import iter_spent_rewards
+from hathor.transaction import BaseTransaction, TxInput, TxOutput, TxVersion
 from hathor.transaction.base_transaction import TX_HASH_SIZE
 from hathor.transaction.exceptions import InvalidToken
 from hathor.transaction.util import VerboseCallback, unpack, unpack_len
@@ -135,7 +136,7 @@ def _calculate_inherited_min_height(self) -> int:
     def _calculate_my_min_height(self) -> int:
         """ Calculates min height derived from own spent block rewards"""
         min_height = 0
-        for blk in self.iter_spent_rewards():
+        for blk in iter_spent_rewards(self, not_none(self.storage)):
             min_height = max(min_height, blk.get_height() + self._settings.REWARD_SPEND_MIN_BLOCKS + 1)
         return min_height
 
@@ -346,47 +347,6 @@ def _update_token_info_from_outputs(self, *, token_dict: dict[TokenUid, TokenInf
                     sum_tokens = token_info.amount + tx_output.value
                     token_dict[token_uid] = TokenInfo(sum_tokens, token_info.can_mint, token_info.can_melt)
 
-    def iter_spent_rewards(self) -> Iterator[Block]:
-        """Iterate over all the rewards being spent, assumes tx has been verified."""
-        for input_tx in self.inputs:
-            spent_tx = self.get_spent_tx(input_tx)
-            if spent_tx.is_block:
-                assert isinstance(spent_tx, Block)
-                yield spent_tx
-
-    def is_spent_reward_locked(self) -> bool:
-        """ Check whether any spent reward is currently locked, considering only the block rewards spent by this tx
-        itself, and not the inherited `min_height`"""
-        return self.get_spent_reward_locked_info() is not None
-
-    def get_spent_reward_locked_info(self) -> Optional[RewardLockedInfo]:
-        """Check if any input block reward is locked, returning the locked information if any, or None if they are all
-        unlocked."""
-        for blk in self.iter_spent_rewards():
-            assert blk.hash is not None
-            needed_height = self._spent_reward_needed_height(blk)
-            if needed_height > 0:
-                return RewardLockedInfo(blk.hash, needed_height)
-        return None
-
-    def _spent_reward_needed_height(self, block: Block) -> int:
-        """ Returns height still needed to unlock this `block` reward: 0 means it's unlocked."""
-        import math
-        assert self.storage is not None
-        # omitting timestamp to get the current best block, this will usually hit the cache instead of being slow
-        tips = self.storage.get_best_block_tips()
-        assert len(tips) > 0
-        best_height = math.inf
-        for tip in tips:
-            blk = self.storage.get_transaction(tip)
-            assert isinstance(blk, Block)
-            best_height = min(best_height, blk.get_height())
-        assert isinstance(best_height, int)
-        spent_height = block.get_height()
-        spend_blocks = best_height - spent_height
-        needed_height = self._settings.REWARD_SPEND_MIN_BLOCKS - spend_blocks
-        return max(needed_height, 0)
-
     def is_double_spending(self) -> bool:
         """ Iterate through inputs to check if they were already spent
             Used to prevent users from sending double spending transactions to the network

hathor/verification/transaction_verifier.py

@@ -15,6 +15,7 @@
 from hathor.conf.settings import HathorSettings
 from hathor.daa import DifficultyAdjustmentAlgorithm
 from hathor.profiler import get_cpu_profiler
+from hathor.reward_lock import get_spent_reward_locked_info
 from hathor.transaction import BaseTransaction, Transaction, TxInput
 from hathor.transaction.exceptions import (
     ConflictingInputs,
@@ -36,6 +37,7 @@
 from hathor.transaction.transaction import TokenInfo
 from hathor.transaction.util import get_deposit_amount, get_withdraw_amount
 from hathor.types import TokenUid, VertexId
+from hathor.util import not_none
 
 cpu = get_cpu_profiler()
 
@@ -144,7 +146,7 @@ def verify_script(self, *, tx: Transaction, input_tx: TxInput, spent_tx: BaseTra
     def verify_reward_locked(self, tx: Transaction) -> None:
         """Will raise `RewardLocked` if any reward is spent before the best block height is enough, considering only
         the block rewards spent by this tx itself, and not the inherited `min_height`."""
-        info = tx.get_spent_reward_locked_info()
+        info = get_spent_reward_locked_info(tx, not_none(tx.storage))
         if info is not None:
             raise RewardLocked(f'Reward {info.block_hash.hex()} still needs {info.blocks_needed} to be unlocked.')