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_payments := await self.store.get_pending_payment_count()) != 0:
+                    self.log.warning(f"Pending payments ({pending_payments}), waiting")
                     await asyncio.sleep(60)
                     continue
 
@@ -408,14 +403,16 @@ async def create_payment_loop(self):
                             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 = []
+                        for payment in additions_sub_list:
+                            await self.store.add_payment(
+                                payment["puzzle_hash"],
+                                uint64(payment["amount"]),
+                                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()
@@ -441,23 +438,38 @@ 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 payments record")
+                    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
                 # blockchain_fee = 0.00001 * (10 ** 12) * len(payment_targets)
                 blockchain_fee = 0
                 try:
+                    await self.store.update_is_payment(
+                        [(puzzle_hash, timestamp) for puzzle_hash, _, timestamp, _ in pending_payments], True
+                    )
                     transaction: TransactionRecord = await self.wallet_rpc_client.send_transaction_multi(
                         self.wallet_id, payment_targets, fee=blockchain_fee
                     )
-                except ValueError as e:
+                except Exception as e:
                     self.log.error(f"Error making payment: {e}")
+                    await self.store.update_is_payment(
+                        [(puzzle_hash, timestamp) for puzzle_hash, _, timestamp, _ in pending_payments], False
+                    )
                     await asyncio.sleep(10)
-                    await self.pending_payments.put(payment_targets)
                     continue
 
                 self.log.info(f"Transaction: {transaction}")

pool/store/abstract.py

@@ -68,3 +68,19 @@ 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, puzzle_hash: bytes32, amount: uint64, timestamp: uint64, is_payment: bool):
+        """Persist a new payment record in the store"""
+
+    @abstractmethod
+    async def update_is_payment(self, puzzle_hash_timestamp: List[Tuple[bytes32, uint64]], is_payment: bool):
+        """Update is_payment for payment records identified by ``puzzle_hash`` and ``timestamp``"""
+
+    @abstractmethod
+    async def get_pending_payment_records(self, count: int) -> List[Tuple[bytes32, 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

@@ -48,9 +48,21 @@ 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(
+                    puzzle_hash text,
+                    amount 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 puzzle_hash_index on payment(puzzle_hash)")
 
         await self.connection.commit()
 
@@ -192,3 +204,45 @@ 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, puzzle_hash: bytes32, amount: uint64, timestamp: uint64, is_payment: bool):
+        cursor = await self.connection.execute(
+            "INSERT INTO payment VALUES(?, ?, ?, ?)",
+            (puzzle_hash.hex(), amount, timestamp, int(is_payment)),
+        )
+        await cursor.close()
+        await self.connection.commit()
+
+    async def update_is_payment(self, puzzle_hash_timestamp: List[Tuple[bytes32, uint64]], is_payment: bool):
+        cursor = await self.connection.executemany(
+            "UPDATE payment SET is_payment=? WHERE puzzle_hash=? AND timestamp=?",
+            tuple((int(is_payment), ph.hex(), timestamp) for ph, timestamp in puzzle_hash_timestamp),
+        )
+        await cursor.close()
+        await self.connection.commit()
+
+    async def get_pending_payment_records(self, count: int) -> List[Tuple[bytes32, uint64, uint64, bool]]:
+        cursor = await self.connection.execute(
+            """
+                SELECT
+                    puzzle_hash, amount, 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, uint64, uint64, bool]
+        ] = [
+            (bytes32(bytes.fromhex(puzzle_hash)), uint64(amount), uint64(timestamp), bool(is_payment))
+            for puzzle_hash, amount, 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