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_payment_count_by_is_payment(False)) != 0:
+                    self.log.warning(f"Pending payments ({pending_payment_count}), waiting")
                     await asyncio.sleep(60)
                     continue
 
@@ -408,14 +403,11 @@ 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 ph, amount in additions_sub_list:
+                            await self.store.add_payment(ph, uint64(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,9 +433,24 @@ async def submit_payment_loop(self):
                     await asyncio.sleep(60)
                     continue
 
-                payment_targets = await self.pending_payments.get()
-                assert len(payment_targets) > 0
+                payment_records = await self.store.get_payment_records_by_is_payment(
+                    False, self.max_additions_per_transaction
+                )
+                if len(payment_records) == 0:
+                    self.log.info("No funds to payment record")
+                    await asyncio.sleep(60)
+                    continue
 
+                payment_targets: List[Dict] = [
+                    {
+                        "puzzle_hash": puzzle_hash,
+                        "amount": amount,
+                    }
+                    for puzzle_hash, amount, _, _ in payment_records
+                ]
+                await self.store.update_payments_is_payment(
+                    [(puzzle_hash, timestamp) for puzzle_hash, _, timestamp, _ in payment_records], True
+                )
                 self.log.info(f"Submitting a payment: {payment_targets}")
 
                 # TODO(pool): make sure you have enough to pay the blockchain fee, this will be taken out of the pool
@@ -457,7 +464,9 @@ 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)
+                    await self.store.update_payments_is_payment(
+                        [(puzzle_hash, timestamp) for puzzle_hash, _, timestamp, _ in payment_records], False
+                    )
                     continue
 
                 self.log.info(f"Transaction: {transaction}")

pool/store/abstract.py

@@ -68,3 +68,21 @@ 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_payments_is_payment(self, puzzle_hash_timestamp: List[Tuple[bytes32, uint64]], is_payment: bool):
+        """Update payments is_payment"""
+
+    @abstractmethod
+    async def get_payment_records_by_is_payment(self, is_payment: bool, count: int) -> List[
+        Tuple[bytes32, uint64, uint64, bool]
+    ]:
+        """Fetch payments matching given is_payment"""
+
+    @abstractmethod
+    async def get_payment_count_by_is_payment(self, is_payment: bool) -> int:
+        """Get payment count matching given is_payment"""

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,50 @@ 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_payments_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_payment_records_by_is_payment(self, is_payment: bool, 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=? ORDER BY timestamp ASC LIMIT ?
+            """,
+            (int(is_payment), 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_payment_count_by_is_payment(self, is_payment: bool) -> int:
+        cursor = await self.connection.execute(
+            "SELECT COUNT(*) FROM payment WHERE is_payment=?",
+            (int(is_payment),),
+        )
+        count: int = (await cursor.fetchone())[0]
+        await cursor.close()
+        return count