From 12304da15d0af4b902ed80456e451784ab5fc36c Mon Sep 17 00:00:00 2001
From: Jan Segre <jan@hathor.network>
Date: Tue, 19 Mar 2024 23:23:28 -0300
Subject: [PATCH] fix(indexes): efficient rocksdb mempool-tips initialization

---
 hathor/indexes/manager.py            |  5 ++--
 hathor/indexes/mempool_tips_index.py | 42 ++++++++++++++++++++++------
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/hathor/indexes/manager.py b/hathor/indexes/manager.py
index e38e3ad85..d67aa3776 100644
--- a/hathor/indexes/manager.py
+++ b/hathor/indexes/manager.py
@@ -348,7 +348,6 @@ def enable_utxo_index(self) -> None:
             self.utxo = RocksDBUtxoIndex(self._db)
 
     def enable_mempool_index(self) -> None:
-        from hathor.indexes.memory_mempool_tips_index import MemoryMempoolTipsIndex
+        from hathor.indexes.rocksdb_mempool_tips_index import RocksDBMempoolTipsIndex
         if self.mempool_tips is None:
-            # XXX: use of RocksDBMempoolTipsIndex is very slow and was suspended
-            self.mempool_tips = MemoryMempoolTipsIndex()
+            self.mempool_tips = RocksDBMempoolTipsIndex(self._db)
diff --git a/hathor/indexes/mempool_tips_index.py b/hathor/indexes/mempool_tips_index.py
index 222cc8140..290b4f865 100644
--- a/hathor/indexes/mempool_tips_index.py
+++ b/hathor/indexes/mempool_tips_index.py
@@ -14,6 +14,7 @@
 
 from abc import abstractmethod
 from collections.abc import Collection
+from itertools import chain
 from typing import TYPE_CHECKING, Iterable, Iterator, Optional, cast
 
 import structlog
@@ -27,22 +28,33 @@
     from hathor.transaction.storage import TransactionStorage
 
 SCOPE = Scope(
-    include_blocks=True,
+    include_blocks=False,
     include_txs=True,
     include_voided=True,
+    topological_order=False,
 )
 
 
+def any_non_voided(tx_storage: 'TransactionStorage', hashes: Iterable[bytes]) -> bool:
+    """
+    If there's any vertex with hash in hashes that is not voided, this function returns True, otherwise False.
+
+    Notice that this means that an empty `hashes` also returns False.
+    """
+    for tx_hash in hashes:
+        tx = tx_storage.get_transaction(tx_hash)
+        tx_meta = tx.get_metadata()
+        if not tx_meta.voided_by:
+            return True
+    return False
+
+
 class MempoolTipsIndex(BaseIndex):
     """Index to access the tips of the mempool transactions, which haven't been confirmed by a block."""
 
     def get_scope(self) -> Scope:
         return SCOPE
 
-    def init_loop_step(self, tx: BaseTransaction) -> None:
-        self.update(tx)
-
-    # originally tx_storage.update_mempool_tips
     @abstractmethod
     def update(self, tx: BaseTransaction, *, remove: Optional[bool] = None) -> None:
         """
@@ -52,7 +64,6 @@ def update(self, tx: BaseTransaction, *, remove: Optional[bool] = None) -> None:
         """
         raise NotImplementedError
 
-    # originally tx_storage.iter_mempool_tips
     @abstractmethod
     def iter(self, tx_storage: 'TransactionStorage', max_timestamp: Optional[float] = None) -> Iterator[Transaction]:
         """
@@ -60,7 +71,6 @@ def iter(self, tx_storage: 'TransactionStorage', max_timestamp: Optional[float]
         """
         raise NotImplementedError
 
-    # originally tx_storage.iter_mempool
     @abstractmethod
     def iter_all(self, tx_storage: 'TransactionStorage') -> Iterator[Transaction]:
         """
@@ -68,7 +78,6 @@ def iter_all(self, tx_storage: 'TransactionStorage') -> Iterator[Transaction]:
         """
         raise NotImplementedError
 
-    # originally tx_storage.get_mempool_tips_index
     @abstractmethod
     def get(self) -> set[bytes]:
         """
@@ -85,6 +94,23 @@ class ByteCollectionMempoolTipsIndex(MempoolTipsIndex):
     log: 'structlog.stdlib.BoundLogger'
     _index: 'Collection[bytes]'
 
+    def init_loop_step(self, tx: BaseTransaction) -> None:
+        assert tx.hash is not None
+        assert tx.storage is not None
+        tx_meta = tx.get_metadata()
+        # do not include transactions that have been confirmed
+        if tx_meta.first_block:
+            return
+        tx_storage = tx.storage
+        # do not include transactions that have a non-voided child
+        if any_non_voided(tx_storage, tx_meta.children):
+            return
+        # do not include transactions that have a non-voided spent output
+        if any_non_voided(tx_storage, chain(*tx_meta.spent_outputs.values())):
+            return
+        # include them otherwise
+        self._add(tx.hash)
+
     @abstractmethod
     def _discard(self, tx: bytes) -> None:
         raise NotImplementedError()