pool: Improve transaction (Chia-Network#201)

Author: jianghushinian

pool/pool.py

@@ -155,9 +155,6 @@ def __init__(
         # faster.
         self.max_additions_per_transaction = pool_config["max_additions_per_transaction"]
 
-        # This is the list of payments that we have not sent yet, to farmers
-        self.pending_payments: Optional[asyncio.Queue] = None
-
         # Keeps track of the latest state of our node
         self.blockchain_state = {"peak": None}
 
@@ -207,8 +204,6 @@ async def start(self):
         self.submit_payment_loop_task = asyncio.create_task(self.submit_payment_loop())
         self.get_peak_loop_task = asyncio.create_task(self.get_peak_loop())
 
-        self.pending_payments = asyncio.Queue()
-
     async def stop(self):
         if self.confirm_partials_loop_task is not None:
             self.confirm_partials_loop_task.cancel()
@@ -362,8 +357,8 @@ async def create_payment_loop(self):
                     await asyncio.sleep(60)
                     continue
 
-                if self.pending_payments.qsize() != 0:
-                    self.log.warning(f"Pending payments ({self.pending_payments.qsize()}), waiting")
+                if (pending_payment_count := await self.store.get_pending_payment_count()) != 0:
+                    self.log.warning(f"Pending payments ({pending_payment_count}), waiting")
                     await asyncio.sleep(60)
                     continue
 
@@ -391,34 +386,46 @@ async def create_payment_loop(self):
                 self.log.info(f"Total amount to distribute: {amount_to_distribute  / (10 ** 12)}")
 
                 async with self.store.lock:
-                    # Get the points of each farmer, as well as payout instructions. Here a chia address is used,
-                    # but other blockchain addresses can also be used.
-                    points_and_ph: List[
-                        Tuple[uint64, bytes]
-                    ] = await self.store.get_farmer_points_and_payout_instructions()
-                    total_points = sum([pt for (pt, ph) in points_and_ph])
+                    # Get the launcher_id and points of each farmer, as well as payout instructions.
+                    # Here a chia address is used, but other blockchain addresses can also be used.
+                    launcher_id_and_points_and_ph: List[
+                        Tuple[bytes32, uint64, bytes32]
+                    ] = await self.store.get_farmer_launcher_id_and_points_and_payout_instructions()
+                    total_points = sum([pt for (launcher_id, pt, ph) in launcher_id_and_points_and_ph])
                     if total_points > 0:
                         mojo_per_point = floor(amount_to_distribute / total_points)
                         self.log.info(f"Paying out {mojo_per_point} mojo / point")
 
+                        # Pool fee payment record launcher_id is equal to puzzle_hash, points is equal to 0.
                         additions_sub_list: List[Dict] = [
-                            {"puzzle_hash": self.pool_fee_puzzle_hash, "amount": pool_coin_amount}
+                            {
+                                "launcher_id": self.pool_fee_puzzle_hash,
+                                "puzzle_hash": self.pool_fee_puzzle_hash,
+                                "amount": pool_coin_amount,
+                                "points": 0,
+                            }
                         ]
-                        for points, ph in points_and_ph:
+                        for launcher_id, points, ph in launcher_id_and_points_and_ph:
                             if points > 0:
-                                additions_sub_list.append({"puzzle_hash": ph, "amount": points * mojo_per_point})
-
-                            if len(additions_sub_list) == self.max_additions_per_transaction:
-                                await self.pending_payments.put(additions_sub_list.copy())
-                                self.log.info(f"Will make payments: {additions_sub_list}")
-                                additions_sub_list = []
+                                additions_sub_list.append({
+                                    "launcher_id": launcher_id,
+                                    "puzzle_hash": ph,
+                                    "amount": points * mojo_per_point,
+                                    "points": points,
+                                })
+
+                        for payment in additions_sub_list:
+                            await self.store.add_payment(
+                                payment["launcher_id"],
+                                payment["puzzle_hash"],
+                                uint64(payment["amount"]),
+                                payment["points"],
+                                uint64(int(time.time())),
+                                False,
+                            )
 
                         if len(additions_sub_list) > 0:
                             self.log.info(f"Will make payments: {additions_sub_list}")
-                            await self.pending_payments.put(additions_sub_list.copy())
-
-                        # Subtract the points from each farmer
-                        await self.store.clear_farmer_points()
                     else:
                         self.log.info(f"No points for any farmer. Waiting {self.payment_interval}")
 
@@ -441,10 +448,20 @@ async def submit_payment_loop(self):
                     await asyncio.sleep(60)
                     continue
 
-                payment_targets = await self.pending_payments.get()
-                assert len(payment_targets) > 0
+                pending_payments = await self.store.get_pending_payment_records(self.max_additions_per_transaction)
+                if len(pending_payments) == 0:
+                    self.log.info("No funds to pending payment records")
+                    await asyncio.sleep(60)
+                    continue
+                self.log.info(f"Submitting a payment: {pending_payments}")
 
-                self.log.info(f"Submitting a payment: {payment_targets}")
+                payment_targets: List[Dict] = [
+                    {
+                        "puzzle_hash": puzzle_hash,
+                        "amount": amount,
+                    }
+                    for _, puzzle_hash, amount, _, _, _ in pending_payments
+                ]
 
                 # TODO(pool): make sure you have enough to pay the blockchain fee, this will be taken out of the pool
                 # fee itself. Alternatively you can set it to 0 and wait longer
@@ -457,34 +474,40 @@ async def submit_payment_loop(self):
                 except ValueError as e:
                     self.log.error(f"Error making payment: {e}")
                     await asyncio.sleep(10)
-                    await self.pending_payments.put(payment_targets)
                     continue
 
                 self.log.info(f"Transaction: {transaction}")
 
-                while (
-                    not transaction.confirmed
-                    or not (peak_height - transaction.confirmed_at_height) > self.confirmation_security_threshold
-                ):
-                    transaction = await self.wallet_rpc_client.get_transaction(self.wallet_id, transaction.name)
-                    peak_height = self.blockchain_state["peak"].height
-                    self.log.info(
-                        f"Waiting for transaction to obtain {self.confirmation_security_threshold} confirmations"
+                async with self.store.tx():
+                    await self.store.update_is_payment(
+                        [
+                            (launcher_id, points, timestamp)
+                            for launcher_id, _, _, points, timestamp, _ in pending_payments
+                        ],
+                        is_payment=True,
+                        auto_commit=False,
                     )
-                    if not transaction.confirmed:
-                        self.log.info(f"Not confirmed. In mempool? {transaction.is_in_mempool()}")
-                    else:
-                        self.log.info(f"Confirmations: {peak_height - transaction.confirmed_at_height}")
-                    await asyncio.sleep(10)
 
-                # TODO(pool): persist in DB
-                self.log.info(f"Successfully confirmed payments {payment_targets}")
+                    while (
+                        not transaction.confirmed
+                        or not (peak_height - transaction.confirmed_at_height) > self.confirmation_security_threshold
+                    ):
+                        transaction = await self.wallet_rpc_client.get_transaction(self.wallet_id, transaction.name)
+                        peak_height = self.blockchain_state["peak"].height
+                        self.log.info(
+                            f"Waiting for transaction to obtain {self.confirmation_security_threshold} confirmations"
+                        )
+                        if not transaction.confirmed:
+                            self.log.info(f"Not confirmed. In mempool? {transaction.is_in_mempool()}")
+                        else:
+                            self.log.info(f"Confirmations: {peak_height - transaction.confirmed_at_height}")
+                        await asyncio.sleep(10)
 
+                self.log.info(f"Successfully confirmed payments {pending_payments}")
             except asyncio.CancelledError:
                 self.log.info("Cancelled submit_payment_loop, closing")
                 return
             except Exception as e:
-                # TODO(pool): retry transaction if failed
                 self.log.error(f"Unexpected error in submit_payment_loop: {e}")
                 await asyncio.sleep(60)
 

pool/store/abstract.py

@@ -23,6 +23,10 @@ def __init__(self):
     async def connect(self):
         """Perform IO-related initialization"""
 
+    @abstractmethod
+    async def tx(self):
+        """Performing Transactions for async with statement"""
+
     @abstractmethod
     async def add_farmer_record(self, farmer_record: FarmerRecord, metadata: RequestMetadata):
         """Persist a new Farmer in the store"""
@@ -54,7 +58,7 @@ async def get_farmer_records_for_p2_singleton_phs(self, puzzle_hashes: Set[bytes
         """Fetch Farmers matching given puzzle hashes"""
 
     @abstractmethod
-    async def get_farmer_points_and_payout_instructions(self) -> List[Tuple[uint64, bytes]]:
+    async def get_farmer_launcher_id_and_points_and_payout_instructions(self) -> List[Tuple[bytes32, uint64, bytes32]]:
         """Fetch all farmers and their respective payout instructions"""
 
     @abstractmethod
@@ -68,3 +72,38 @@ async def add_partial(self, launcher_id: bytes32, timestamp: uint64, difficulty:
     @abstractmethod
     async def get_recent_partials(self, launcher_id: bytes32, count: int) -> List[Tuple[uint64, uint64]]:
         """Fetch last ``count`` partials for Farmer identified by ``launcher_id``"""
+
+    @abstractmethod
+    async def add_payment(
+        self,
+        launcher_id: bytes32,
+        puzzle_hash: bytes32,
+        amount: uint64,
+        points: int,
+        timestamp: uint64,
+        is_payment: bool,
+    ):
+        """Persist a new payment record in the store"""
+
+    @abstractmethod
+    async def update_is_payment(
+        self,
+        launcher_id_and_points_and_timestamp: List[Tuple[bytes32, uint64, uint64]],
+        is_payment: bool,
+        auto_commit: bool = True,
+    ):
+        """
+        Update is_payment for payment records identified by ``launcher_id`` and ``timestamp``.
+        When is_payment is True, the payouts is completed, subtract the points from farmer.
+        auto_commit decides whether to commit the transaction.
+        """
+
+    @abstractmethod
+    async def get_pending_payment_records(self, count: int) -> List[
+        Tuple[bytes32, bytes32, uint64, uint64, uint64, bool]
+    ]:
+        """Fetch ``count`` pending payment records"""
+
+    @abstractmethod
+    async def get_pending_payment_count(self) -> int:
+        """Fetch pending payment records count"""

pool/store/sqlite_store.py

@@ -1,3 +1,4 @@
+from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Optional, Set, List, Tuple, Dict
 
@@ -48,9 +49,23 @@ async def connect(self):
             "CREATE TABLE IF NOT EXISTS partial(launcher_id text, timestamp bigint, difficulty bigint)"
         )
 
+        await self.connection.execute(
+            """
+                CREATE TABLE IF NOT EXISTS payment(
+                    launcher_id text,
+                    puzzle_hash text,
+                    amount bigint,
+                    points bigint,
+                    timestamp bigint,
+                    is_payment tinyint
+                )
+            """
+        )
+
         await self.connection.execute("CREATE INDEX IF NOT EXISTS scan_ph on farmer(p2_singleton_puzzle_hash)")
         await self.connection.execute("CREATE INDEX IF NOT EXISTS timestamp_index on partial(timestamp)")
         await self.connection.execute("CREATE INDEX IF NOT EXISTS launcher_id_index on partial(launcher_id)")
+        await self.connection.execute("CREATE INDEX IF NOT EXISTS launcher_id_index on payment(launcher_id)")
 
         await self.connection.commit()
 
@@ -70,6 +85,15 @@ def _row_to_farmer_record(row) -> FarmerRecord:
             True if row[10] == 1 else False,
         )
 
+    @asynccontextmanager
+    async def tx(self):
+        try:
+            yield
+            await self.connection.commit()
+        except Exception:
+            await self.connection.rollback()
+            raise
+
     async def add_farmer_record(self, farmer_record: FarmerRecord, metadata: RequestMetadata):
         cursor = await self.connection.execute(
             f"INSERT OR REPLACE INTO farmer VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
@@ -146,21 +170,22 @@ async def get_farmer_records_for_p2_singleton_phs(self, puzzle_hashes: Set[bytes
         rows = await cursor.fetchall()
         return [self._row_to_farmer_record(row) for row in rows]
 
-    async def get_farmer_points_and_payout_instructions(self) -> List[Tuple[uint64, bytes]]:
-        cursor = await self.connection.execute(f"SELECT points, payout_instructions from farmer")
+    async def get_farmer_launcher_id_and_points_and_payout_instructions(self) -> List[Tuple[bytes32, uint64, bytes32]]:
+        cursor = await self.connection.execute(f"SELECT launcher_id, points, payout_instructions from farmer")
         rows = await cursor.fetchall()
-        accumulated: Dict[bytes32, uint64] = {}
+        accumulated: Dict[Tuple[bytes32, bytes32], uint64] = {}
         for row in rows:
-            points: uint64 = uint64(row[0])
-            ph: bytes32 = bytes32(bytes.fromhex(row[1]))
+            launcher_id: bytes32 = bytes32(bytes.fromhex(row[0]))
+            points: uint64 = uint64(row[1])
+            ph: bytes32 = bytes32(bytes.fromhex(row[2]))
             if ph in accumulated:
-                accumulated[ph] += points
+                accumulated[(launcher_id, ph)] += points
             else:
-                accumulated[ph] = points
+                accumulated[(launcher_id, ph)] = points
 
-        ret: List[Tuple[uint64, bytes32]] = []
-        for ph, total_points in accumulated.items():
-            ret.append((total_points, ph))
+        ret: List[Tuple[bytes32, uint64, bytes32]] = []
+        for (launcher_id, ph), total_points in accumulated.items():
+            ret.append((launcher_id, total_points, ph))
         return ret
 
     async def clear_farmer_points(self) -> None:
@@ -192,3 +217,81 @@ async def get_recent_partials(self, launcher_id: bytes32, count: int) -> List[Tu
         rows = await cursor.fetchall()
         ret: List[Tuple[uint64, uint64]] = [(uint64(timestamp), uint64(difficulty)) for timestamp, difficulty in rows]
         return ret
+
+    async def add_payment(
+        self,
+        launcher_id: bytes32,
+        puzzle_hash: bytes32,
+        amount: uint64,
+        points: int,
+        timestamp: uint64,
+        is_payment: bool,
+    ):
+        cursor = await self.connection.execute(
+            "INSERT INTO payment VALUES(?, ?, ?, ?, ?, ?)",
+            (launcher_id.hex(), puzzle_hash.hex(), amount, points, timestamp, int(is_payment)),
+        )
+        await cursor.close()
+        await self.connection.commit()
+
+    async def update_is_payment(
+        self,
+        launcher_id_and_points_and_timestamp: List[Tuple[bytes32, uint64, uint64]],
+        is_payment: bool,
+        auto_commit: bool = True,
+    ):
+        cursor = await self.connection.executemany(
+            "UPDATE payment SET is_payment=? WHERE launcher_id=? AND timestamp=?",
+            tuple(
+                (int(is_payment), launcher_id.hex(), timestamp)
+                for launcher_id, _, timestamp in launcher_id_and_points_and_timestamp
+            ),
+        )
+        await cursor.close()
+
+        if is_payment is True:
+            cursor = await self.connection.executemany(
+                "UPDATE farmer SET points=points-? WHERE launcher_id=?",
+                tuple(
+                    (points, launcher_id.hex())
+                    for launcher_id, points, _ in launcher_id_and_points_and_timestamp
+                ),
+            )
+            await cursor.close()
+
+        if auto_commit is True:
+            await self.connection.commit()
+
+    async def get_pending_payment_records(self, count: int) -> List[
+        Tuple[bytes32, bytes32, uint64, uint64, uint64, bool]
+    ]:
+        cursor = await self.connection.execute(
+            """
+                SELECT
+                    launcher_id, puzzle_hash, amount, points, timestamp, is_payment
+                FROM payment
+                WHERE is_payment=0 ORDER BY timestamp ASC LIMIT ?
+            """,
+            (count,),
+        )
+        rows = await cursor.fetchall()
+        await cursor.close()
+        ret: List[
+            Tuple[bytes32, bytes32, uint64, uint64, uint64, bool]
+        ] = [
+            (
+                bytes32(bytes.fromhex(launcher_id)),
+                bytes32(bytes.fromhex(puzzle_hash)),
+                uint64(amount),
+                uint64(points),
+                uint64(timestamp),
+                bool(is_payment),
+            ) for launcher_id, puzzle_hash, amount, points, timestamp, is_payment in rows
+        ]
+        return ret
+
+    async def get_pending_payment_count(self) -> int:
+        cursor = await self.connection.execute("SELECT COUNT(*) FROM payment WHERE is_payment=0")
+        count: int = (await cursor.fetchone())[0]
+        await cursor.close()
+        return count