txbatcher: use lock, rename private methods, add type hints

Author: ecdsa

electrum/txbatcher.py

@@ -1,6 +1,8 @@
 import asyncio
+import threading
 import copy
 
+from typing import Dict, Sequence
 from . import util
 from .logging import Logger
 from .util import log_exceptions
@@ -74,6 +76,8 @@ def __init__(self, wallet):
         self.wallet = wallet
         self.config = wallet.config
         self.db = wallet.db
+        self.lock = threading.Lock()
+        # fixme: not robust to client restart, because we do not persist batch_payments
         self.batch_payments = []       # list of payments we need to make
         self.batch_inputs = {}         # list of inputs we need to sweep
 
@@ -96,7 +100,8 @@ def __init__(self, wallet):
 
     def add_batch_payment(self, output: 'PartialTxOutput'):
         # todo: maybe we should raise NotEnoughFunds here
-        self.batch_payments.append(output)
+        with self.lock:
+            self.batch_payments.append(output)
 
     def add_sweep_info(self, sweep_info: 'SweepInfo'):
         txin = sweep_info.txin
@@ -123,7 +128,6 @@ def add_sweep_info(self, sweep_info: 'SweepInfo'):
                         base_txin.witness_script = txin.witness_script
                         base_txin.script_sig = txin.script_sig
 
-
     def get_base_tx(self) -> Optional[Transaction]:
         if self._base_tx:
             return self._base_tx
@@ -140,7 +144,7 @@ def get_base_tx(self) -> Optional[Transaction]:
         base_tx.add_info_from_wallet(self.wallet) # needed for txid
         return base_tx
 
-    def find_confirmed_base_tx(self) -> Optional[Transaction]:
+    def _find_confirmed_base_tx(self) -> Optional[Transaction]:
         for txid in self.batch_txids:
             tx_mined_status = self.wallet.adb.get_tx_height(txid)
             if tx_mined_status.conf > 0:
@@ -149,8 +153,7 @@ def find_confirmed_base_tx(self) -> Optional[Transaction]:
                 tx.add_info_from_wallet(self.wallet) # needed for txid
                 return tx
 
-    def to_pay_after(self, tx):
-        # fixme: not robust to client restart, because we do not persist batch_payments
+    def _to_pay_after(self, tx) -> Sequence[PartialTxOutput]:
         if not tx:
             return self.batch_payments
         to_pay = []
@@ -162,7 +165,7 @@ def to_pay_after(self, tx):
                 outputs.remove(x)
         return to_pay
 
-    def to_sweep_after(self, tx):
+    def _to_sweep_after(self, tx) -> Dict[str, SweepInfo]:
         tx_prevouts = set(txin.prevout for txin in tx.inputs()) if tx else set()
         result = []
         for k,v in self.batch_inputs.items():
@@ -179,7 +182,7 @@ def to_sweep_after(self, tx):
             result.append((k,v))
         return dict(result)
 
-    def should_bump_fee(self, base_tx):
+    def _should_bump_fee(self, base_tx) -> bool:
         if base_tx is None:
             return False
         base_tx_fee = base_tx.get_fee()
@@ -196,30 +199,30 @@ async def run(self):
             password = self.wallet.get_unlocked_password()
             if self.wallet.has_keystore_encryption() and not password:
                 continue
-            await self.maybe_broadcast_legacy_htlc_txs()
-            tx = self.find_confirmed_base_tx()
+            await self._maybe_broadcast_legacy_htlc_txs()
+            tx = self._find_confirmed_base_tx()
             if tx:
                 self.logger.info(f'base tx confirmed {tx.txid()}')
-                self.clear_batch_processing(tx)
-                self.start_new_batch(tx)
+                self._clear_batch_processing(tx)
+                self._start_new_batch(tx)
             base_tx = self.get_base_tx()
-            to_pay = self.to_pay_after(base_tx)
-            to_sweep = self.to_sweep_after(base_tx)
+            to_pay = self._to_pay_after(base_tx)
+            to_sweep = self._to_sweep_after(base_tx)
             to_sweep_now = {}
             for k, v in to_sweep.items():
-                can_broadcast, wanted_height = self.can_broadcast(v, base_tx)
+                can_broadcast, wanted_height = self._can_broadcast(v, base_tx)
                 if can_broadcast:
                     to_sweep_now[k] = v
                 else:
                     self.wallet.add_future_tx(v, wanted_height)
-            if not to_pay and not to_sweep_now and not self.should_bump_fee(base_tx):
+            if not to_pay and not to_sweep_now and not self._should_bump_fee(base_tx):
                 continue
             try:
-                tx = self.create_batch_tx(base_tx, to_sweep_now, to_pay, password)
+                tx = self._create_batch_tx(base_tx, to_sweep_now, to_pay, password)
             except Exception as e:
                 self.logger.exception(f'Cannot create batch transaction: {repr(e)}')
                 if base_tx:
-                    self.start_new_batch(base_tx)
+                    self._start_new_batch(base_tx)
                     continue
                 await asyncio.sleep(self.RETRY_DELAY)
                 continue
@@ -242,7 +245,7 @@ async def run(self):
                     self.logger.info(f'starting new batch because could not broadcast')
                     self.start_new_batch(base_tx)
 
-    def create_batch_tx(self, base_tx, to_sweep, to_pay, password):
+    def _create_batch_tx(self, base_tx, to_sweep, to_pay, password):
         self.logger.info(f'to_sweep: {list(to_sweep.keys())}')
         self.logger.info(f'to_pay: {to_pay}')
         inputs = []
@@ -260,7 +263,6 @@ def create_batch_tx(self, base_tx, to_sweep, to_pay, password):
         self.logger.info(f'locktime: {locktime}')
         outputs += to_pay
         inputs += self.get_change_inputs(self._parent_tx) if self._parent_tx else []
-
         tx = self.wallet.create_transaction(
             base_tx=base_tx,
             inputs=inputs,
@@ -275,18 +277,18 @@ def create_batch_tx(self, base_tx, to_sweep, to_pay, password):
         assert tx.is_complete()
         return tx
 
-    def clear_batch_processing(self, tx):
+    def _clear_batch_processing(self, tx):
         # this ensure that we can accept an input again
         # if the spending tx is removed from the blockchain
         # fixme: what if there are several batches?
         for txin in tx.inputs():
             if txin.prevout in self.batch_processing:
                 self.batch_processing.remove(txin.prevout)
 
-    def start_new_batch(self, tx):
+    def _start_new_batch(self, tx):
         use_change = tx and tx.has_change() and any([txout in self.batch_payments for txout in tx.outputs()])
-        self.batch_payments = self.to_pay_after(tx)
-        self.batch_inputs = self.to_sweep_after(tx)
+        self.batch_payments = self._to_pay_after(tx)
+        self.batch_inputs = self._to_sweep_after(tx)
         self.batch_txids.clear()
         self._base_tx = None
         self._parent_tx = tx if use_change else None
@@ -301,7 +303,7 @@ def get_change_inputs(self, parent_tx):
             txin.nsequence = 0xffffffff - 2
         return inputs
 
-    def can_broadcast(self, sweep_info: 'SweepInfo', base_tx):
+    def _can_broadcast(self, sweep_info: 'SweepInfo', base_tx):
         prevout = sweep_info.txin.prevout.to_str()
         name = sweep_info.name
         prev_txid, index = prevout.split(':')
@@ -331,7 +333,7 @@ def can_broadcast(self, sweep_info: 'SweepInfo', base_tx):
             wanted_height = prev_height
         return can_broadcast, wanted_height
 
-    async def maybe_broadcast_legacy_htlc_txs(self):
+    async def _maybe_broadcast_legacy_htlc_txs(self):
         """ pre-anchor htlc txs cannot be batched """
         for sweep_info in list(self.batch_inputs.values()):
             if sweep_info.name == 'first-stage-htlc':

tests/test_txbatcher.py

@@ -118,8 +118,6 @@ async def test_batch_payments(self, mock_save_db):
         assert wallet.adb.get_transaction(tx1.txid()) is not None
         assert wallet.adb.get_transaction(tx1_prime.txid()) is None
         # txbatcher creates tx2
-        self.logger.info(f'to pay after {wallet.txbatcher.to_pay_after(tx1)}')
-        self.logger.info(f'{tx_mined_status}')
         await self.network._tx_event.wait()
         tx2 = wallet.txbatcher.get_base_tx()
         assert output1 in tx1.outputs()