address #275 (comment)

Author: jansegre

hathor/p2p/sync_v2/agent.py

@@ -53,6 +53,10 @@ class PeerState(Enum):
     SYNCING_MEMPOOL = 'syncing-mempool'
 
 
+class _GetDataOrigin(Enum):
+    MEMPOOL = 'mempool'
+
+
 class NodeBlockSync(SyncAgent):
     """ An algorithm to sync two peers based on their blockchain.
     """
@@ -1000,33 +1004,34 @@ def handle_transaction(self, payload: str) -> None:
                 self.log.debug('tx streaming in progress', txs_received=self._tx_received)
 
     @inlineCallbacks
-    def get_tx(self, tx_id: bytes) -> Generator[Deferred, Any, BaseTransaction]:
+    def get_tx_mempool(self, tx_id: bytes) -> Generator[Deferred, Any, BaseTransaction]:
         """ Async method to get a transaction from the db/cache or to download it.
         """
+        assert self.state is PeerState.SYNCING_MEMPOOL, 'get_tx_mempool must only be called on mempool state'
         tx = self._get_tx_cache.get(tx_id)
         if tx is not None:
             self.log.debug('tx in cache', tx=tx_id.hex())
             return tx
         try:
             tx = self.tx_storage.get_transaction(tx_id)
         except TransactionDoesNotExist:
-            tx = yield self.get_data(tx_id, 'mempool')
+            tx = yield self.get_data(tx_id, _GetDataOrigin.MEMPOOL)
             assert tx is not None
             if tx.hash != tx_id:
                 self.protocol.send_error_and_close_connection(f'DATA mempool {tx_id.hex()} hash mismatch')
                 raise
         return tx
 
-    def get_data(self, tx_id: bytes, origin: str) -> Deferred[BaseTransaction]:
+    def get_data(self, tx_id: bytes, origin: _GetDataOrigin) -> Deferred[BaseTransaction]:
         """ Async method to request a tx by id.
         """
         # TODO: deal with stale `get_data` calls
-        if origin != 'mempool':
+        if origin is not _GetDataOrigin.MEMPOOL:
             raise ValueError(f'origin={origin} not supported, only origin=mempool is supported')
         deferred = self._deferred_txs.get(tx_id, None)
         if deferred is None:
             deferred = self._deferred_txs[tx_id] = Deferred()
-            self.send_get_data(tx_id, origin=origin)
+            self.send_get_data(tx_id, origin=origin.name)
             self.log.debug('get_data of new tx_id', deferred=deferred, key=tx_id.hex())
         else:
             # XXX: can we re-use deferred objects like this?

hathor/p2p/sync_v2/mempool.py

@@ -79,7 +79,7 @@ def _unsafe_run(self) -> Generator[Deferred, Any, None]:
         while self.missing_tips:
             self.log.debug('We have missing tips! Let\'s start!', missing_tips=[x.hex() for x in self.missing_tips])
             tx_id = next(iter(self.missing_tips))
-            tx: BaseTransaction = yield self.sync_agent.get_tx(tx_id)
+            tx: BaseTransaction = yield self.sync_agent.get_tx_mempool(tx_id)
             # Stack used by the DFS in the dependencies.
             # We use a deque for performance reasons.
             self.log.debug('start mempool DSF', tx=tx.hash_hex)
@@ -98,7 +98,7 @@ def _dfs(self, stack: deque[BaseTransaction]) -> Generator[Deferred, Any, None]:
                 assert tx == stack.pop()
             else:
                 self.log.debug('Iterate in the DFS.', missing_dep=missing_dep.hex())
-                tx_dep = yield self.sync_agent.get_tx(missing_dep)
+                tx_dep = yield self.sync_agent.get_tx_mempool(missing_dep)
                 stack.append(tx_dep)
                 if len(stack) > self.MAX_STACK_LENGTH:
                     stack.popleft()