Factoring out store

README.md

@@ -9,6 +9,8 @@ Several things are customizable in this pool reference. This includes:
 * Fees to take, and how much to pay in blockchain fees  
 * How farmers' points are counted when paying (PPS, PPLNS, etc)
 * How farmers receive payouts (XCH, BTC, ETH, etc), and how often
+* What store (DB) is used - by default it's an SQLite db. Users can use their own store implementations, based on 
+  `AbstractPoolStore`, by supplying them to `pool_server.start_pool_server`
 
 However, some things cannot be changed. These are described in SPECIFICATION.md, and mostly relate to validation,
 protocol, and the singleton format for smart coins. 
@@ -100,7 +102,7 @@ cd pool-reference
 python3 -m venv ./venv
 source ./venv/bin/activate
 pip install ../chia-blockchain/ 
-sudo CHIA_ROOT="/your/home/dir/.chia/testnet7" ./venv/bin/python pool/pool_server.py
+sudo CHIA_ROOT="/your/home/dir/.chia/testnet7" ./venv/bin/python -m pool
 ```
 
 You should see something like this when starting, but no errors:

pool/__main__.py

@@ -0,0 +1,2 @@
+from pool.pool_server import main
+main()

pool/pool.py

@@ -43,14 +43,17 @@
     launcher_id_to_p2_puzzle_hash,
 )
 
-from difficulty_adjustment import get_new_difficulty
-from singleton import create_absorb_transaction, get_and_validate_singleton_state_inner, get_coin_spend
-from store import FarmerRecord, PoolStore
-from util import error_dict
+from .difficulty_adjustment import get_new_difficulty
+from .singleton import create_absorb_transaction, get_and_validate_singleton_state_inner, get_coin_spend
+from .store.abstract import AbstractPoolStore
+from .store.sqlite_store import SqlitePoolStore
+from .record import FarmerRecord
+from .util import error_dict
 
 
 class Pool:
-    def __init__(self, private_key: PrivateKey, config: Dict, constants: ConsensusConstants):
+    def __init__(self, private_key: PrivateKey, config: Dict, constants: ConsensusConstants,
+                 pool_store: Optional[AbstractPoolStore] = None):
         self.follow_singleton_tasks: Dict[bytes32, asyncio.Task] = {}
         self.log = logging
         # If you want to log to a file: use filename='example.log', encoding='utf-8'
@@ -74,7 +77,7 @@ def __init__(self, private_key: PrivateKey, config: Dict, constants: ConsensusCo
         self.node_rpc_client = None
         self.wallet_rpc_client = None
 
-        self.store: Optional[PoolStore] = None
+        self.store: AbstractPoolStore = pool_store or SqlitePoolStore()
 
         self.pool_fee = pool_config["pool_fee"]
 
@@ -169,7 +172,7 @@ def __init__(self, private_key: PrivateKey, config: Dict, constants: ConsensusCo
         self.wallet_rpc_client: Optional[WalletRpcClient] = None
 
     async def start(self):
-        self.store = await PoolStore.create()
+        await self.store.connect()
         self.pending_point_partials = asyncio.Queue()
 
         self_hostname = self.config["self_hostname"]

pool/pool_server.py

@@ -28,9 +28,10 @@
 from chia.util.default_root import DEFAULT_ROOT_PATH
 from chia.util.config import load_config
 
-from store import FarmerRecord
-from pool import Pool
-from util import error_response
+from .record import FarmerRecord
+from .pool import Pool
+from .store.abstract import AbstractPoolStore
+from .util import error_response
 
 
 def allow_cors(response: web.Response) -> web.Response:
@@ -48,10 +49,11 @@ def check_authentication_token(launcher_id: bytes32, token: uint64, timeout: uin
 
 
 class PoolServer:
-    def __init__(self, private_key: PrivateKey, config: Dict, constants: ConsensusConstants):
+    def __init__(self, private_key: PrivateKey, config: Dict, constants: ConsensusConstants,
+                 pool_store: Optional[AbstractPoolStore] = None):
 
         self.log = logging.getLogger(__name__)
-        self.pool = Pool(private_key, config, constants)
+        self.pool = Pool(private_key, config, constants, pool_store)
 
     async def start(self):
         await self.pool.start()
@@ -247,14 +249,14 @@ async def get_login(self, request_obj) -> web.Response:
 runner = None
 
 
-async def start_pool_server():
+async def start_pool_server(pool_store: Optional[AbstractPoolStore] = None):
     global server
     global runner
     private_key: PrivateKey = AugSchemeMPL.key_gen(std_hash(b"123"))
     config = load_config(DEFAULT_ROOT_PATH, "config.yaml")
     overrides = config["network_overrides"]["constants"][config["selected_network"]]
     constants: ConsensusConstants = DEFAULT_CONSTANTS.replace_str_to_bytes(**overrides)
-    server = PoolServer(private_key, config, constants)
+    server = PoolServer(private_key, config, constants, pool_store)
     await server.start()
 
     # TODO(pool): support TLS
@@ -282,8 +284,12 @@ async def stop():
     await runner.cleanup()
 
 
-if __name__ == "__main__":
+def main():
     try:
         asyncio.run(start_pool_server())
     except KeyboardInterrupt:
         asyncio.run(stop())
+
+
+if __name__ == "__main__":
+    main()

pool/record.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+
+from blspy import G1Element
+from chia.pools.pool_wallet_info import PoolState
+from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.types.coin_solution import CoinSolution
+from chia.util.ints import uint64
+from chia.util.streamable import streamable, Streamable
+
+
+@dataclass(frozen=True)
+@streamable
+class FarmerRecord(Streamable):
+    launcher_id: bytes32  # This uniquely identifies the singleton on the blockchain (ID for this farmer)
+    p2_singleton_puzzle_hash: bytes32  # Derived from the launcher id, delay_time and delay_puzzle_hash
+    delay_time: uint64  # Backup time after which farmer can claim rewards directly, if pool unresponsive
+    delay_puzzle_hash: bytes32  # Backup puzzlehash to claim rewards
+    authentication_public_key: G1Element  # This is the latest public key of the farmer (signs all partials)
+    singleton_tip: CoinSolution  # Last coin solution that is buried in the blockchain, for this singleton
+    singleton_tip_state: PoolState  # Current state of the singleton
+    points: uint64  # Total points accumulated since last rest (or payout)
+    difficulty: uint64  # Current difficulty for this farmer
+    payout_instructions: str  # This is where the pool will pay out rewards to the farmer
+    is_pool_member: bool  # If the farmer leaves the pool, this gets set to False

pool/singleton.py

@@ -17,7 +17,7 @@
 from chia.types.spend_bundle import SpendBundle
 from chia.util.ints import uint32
 
-from store import FarmerRecord
+from .record import FarmerRecord
 
 log = logging
 log.basicConfig(level=logging.INFO)

pool/store/abstract.py

@@ -0,0 +1,68 @@
+from abc import ABC, abstractmethod
+import asyncio
+from typing import Optional, Set, List, Tuple
+
+from chia.pools.pool_wallet_info import PoolState
+from chia.types.blockchain_format.sized_bytes import bytes32
+from chia.types.coin_solution import CoinSolution
+from chia.util.ints import uint64
+
+from ..record import FarmerRecord
+
+
+class AbstractPoolStore(ABC):
+    """
+    Base class for asyncio-related pool stores.
+    """
+    def __init__(self):
+        self.lock = asyncio.Lock()
+
+    @abstractmethod
+    async def connect(self):
+        """Perform IO-related initialization"""
+
+    @abstractmethod
+    async def add_farmer_record(self, farmer_record: FarmerRecord):
+        """Persist a new Farmer in the store"""
+
+    @abstractmethod
+    async def get_farmer_record(self, launcher_id: bytes32) -> Optional[FarmerRecord]:
+        """Fetch a farmer record for given ``launcher_id``. Returns ``None`` if no record found"""
+
+    @abstractmethod
+    async def update_difficulty(self, launcher_id: bytes32, difficulty: uint64):
+        """Update difficulty for Farmer identified by ``launcher_id``"""
+
+    @abstractmethod
+    async def update_singleton(
+        self,
+        launcher_id: bytes32,
+        singleton_tip: CoinSolution,
+        singleton_tip_state: PoolState,
+        is_pool_member: bool,
+    ):
+        """Update Farmer's singleton-related data"""
+
+    @abstractmethod
+    async def get_pay_to_singleton_phs(self) -> Set[bytes32]:
+        """Fetch all puzzle hashes of Farmers in this pool, to scan the blockchain in search of them"""
+
+    @abstractmethod
+    async def get_farmer_records_for_p2_singleton_phs(self, puzzle_hashes: Set[bytes32]) -> List[FarmerRecord]:
+        """Fetch Farmers matching given puzzle hashes"""
+
+    @abstractmethod
+    async def get_farmer_points_and_payout_instructions(self) -> List[Tuple[uint64, bytes]]:
+        """Fetch all farmers and their respective payout instructions"""
+
+    @abstractmethod
+    async def clear_farmer_points(self) -> None:
+        """Rest all Farmers' points to 0"""
+
+    @abstractmethod
+    async def add_partial(self, launcher_id: bytes32, timestamp: uint64, difficulty: uint64):
+        """Register new partial and update corresponding Farmer's points"""
+
+    @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``"""

pool/store.py → pool/store/sqlite_store.py

@@ -1,5 +1,4 @@
 import asyncio
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional, Set, List, Tuple, Dict
 
@@ -10,35 +9,21 @@
 from chia.types.coin_solution import CoinSolution
 from chia.util.ints import uint64
 
-from chia.util.streamable import streamable, Streamable
-
-
-@dataclass(frozen=True)
-@streamable
-class FarmerRecord(Streamable):
-    launcher_id: bytes32  # This uniquely identifies the singleton on the blockchain (ID for this farmer)
-    p2_singleton_puzzle_hash: bytes32  # Derived from the launcher id, delay_time and delay_puzzle_hash
-    delay_time: uint64  # Backup time after which farmer can claim rewards directly, if pool unresponsive
-    delay_puzzle_hash: bytes32  # Backup puzzlehash to claim rewards
-    authentication_public_key: G1Element  # This is the latest public key of the farmer (signs all partials)
-    singleton_tip: CoinSolution  # Last coin solution that is buried in the blockchain, for this singleton
-    singleton_tip_state: PoolState  # Current state of the singleton
-    points: uint64  # Total points accumulated since last rest (or payout)
-    difficulty: uint64  # Current difficulty for this farmer
-    payout_instructions: str  # This is where the pool will pay out rewards to the farmer
-    is_pool_member: bool  # If the farmer leaves the pool, this gets set to False
-
-
-class PoolStore:
-    connection: aiosqlite.Connection
-    lock: asyncio.Lock
-
-    @classmethod
-    async def create(cls):
-        self = cls()
-        self.db_path = Path("pooldb.sqlite")
+from .abstract import AbstractPoolStore
+from ..record import FarmerRecord
+
+
+class SqlitePoolStore(AbstractPoolStore):
+    """
+    Pool store based on SQLite.
+    """
+    def __init__(self, db_path: Path = Path('pooldb.sqlite')):
+        super().__init__()
+        self.db_path = db_path
+        self.connection: Optional[aiosqlite.Connection] = None
+
+    async def connect(self):
         self.connection = await aiosqlite.connect(self.db_path)
-        self.lock = asyncio.Lock()
         await self.connection.execute("pragma journal_mode=wal")
         await self.connection.execute("pragma synchronous=2")
         await self.connection.execute(
@@ -68,8 +53,6 @@ async def create(cls):
 
         await self.connection.commit()
 
-        return self
-
     @staticmethod
     def _row_to_farmer_record(row) -> FarmerRecord:
         return FarmerRecord(