lnchannel/lnhtlc: speed up balance calculation for recent ctns

Move the balance calculation from lnchannel to lnhtlc.
Maintain a running balance in lnhtlc that is coupled with _maybe_active_htlc_ids
for practicality reasons.

electrum/lnchannel.py

@@ -610,7 +610,7 @@ class Channel(Logger):
                 reason = self._receive_fail_reasons.get(htlc.htlc_id)
                 self.lnworker.payment_failed(self, htlc.payment_hash, reason)
 
-    def balance(self, whose, *, ctx_owner=HTLCOwner.LOCAL, ctn=None):
+    def balance(self, whose: HTLCOwner, *, ctx_owner=HTLCOwner.LOCAL, ctn: int = None) -> int:
         """
         This balance in mSAT is not including reserve and fees.
         So a node cannot actually use its whole balance.
@@ -623,22 +623,10 @@ class Channel(Logger):
         """
         assert type(whose) is HTLCOwner
         initial = self.config[whose].initial_msat
-
+        return self.hm.get_balance_msat(whose=whose,
-        # TODO slow. -- and 'balance' is called from a decent number of places (e.g. 'make_commitment')
+                                        ctx_owner=ctx_owner,
-        for direction, htlc in self.hm.all_settled_htlcs_ever(ctx_owner, ctn):
+                                        ctn=ctn,
-            # note: could "simplify" to (whose * ctx_owner == direction * SENT)
+                                        initial_balance_msat=initial)
-            if whose == ctx_owner:
-                if direction == SENT:
-                    initial -= htlc.amount_msat
-                else:
-                    initial += htlc.amount_msat
-            else:
-                if direction == SENT:
-                    initial += htlc.amount_msat
-                else:
-                    initial -= htlc.amount_msat
-
-        return initial
 
     def balance_minus_outgoing_htlcs(self, whose: HTLCOwner, *, ctx_owner: HTLCOwner = HTLCOwner.LOCAL):
         """

electrum/lnhtlc.py

@@ -173,10 +173,12 @@ class HTLCManager:
         self.log['unacked_local_updates2'].pop(self.log[REMOTE]['ctn'], None)
 
     def _update_maybe_active_htlc_ids(self) -> None:
-        # Loosely, we want a set that contains the htlcs that are
+        # - Loosely, we want a set that contains the htlcs that are
-        # not "removed and revoked from all ctxs of both parties".
+        #   not "removed and revoked from all ctxs of both parties". (self._maybe_active_htlc_ids)
-        # It is guaranteed that those htlcs are in the set, but older htlcs might be there too:
+        #   It is guaranteed that those htlcs are in the set, but older htlcs might be there too:
-        # there is a sanity margin of 1 ctn -- this relaxes the care needed re order of method calls.
+        #   there is a sanity margin of 1 ctn -- this relaxes the care needed re order of method calls.
+        # - balance_delta is in sync with maybe_active_htlc_ids. When htlcs are removed from the latter,
+        #   balance_delta is updated to reflect that htlc.
         sanity_margin = 1
         for htlc_proposer in (LOCAL, REMOTE):
             for log_action in ('settles', 'fails'):
@@ -188,10 +190,14 @@ class HTLCManager:
                             and ctns[REMOTE] is not None
                             and ctns[REMOTE] <= self.ctn_oldest_unrevoked(REMOTE) - sanity_margin):
                         self._maybe_active_htlc_ids[htlc_proposer].remove(htlc_id)
+                        if log_action == 'settles':
+                            htlc = self.log[htlc_proposer]['adds'][htlc_id]  # type: UpdateAddHtlc
+                            self._balance_delta -= htlc.amount_msat * htlc_proposer
 
     def _init_maybe_active_htlc_ids(self):
         self._maybe_active_htlc_ids = {LOCAL: set(), REMOTE: set()}  # first idx is "side who offered htlc"
         # add all htlcs
+        self._balance_delta = 0  # the balance delta of LOCAL since channel open
         for htlc_proposer in (LOCAL, REMOTE):
             for htlc_id in self.log[htlc_proposer]['adds']:
                 self._maybe_active_htlc_ids[htlc_proposer].add(htlc_id)
@@ -333,6 +339,37 @@ class HTLCManager:
         received = [(RECEIVED, x) for x in self.all_settled_htlcs_ever_by_direction(subject, RECEIVED, ctn)]
         return sent + received
 
+    def get_balance_msat(self, whose: HTLCOwner, *, ctx_owner=HTLCOwner.LOCAL, ctn: int = None,
+                         initial_balance_msat: int) -> int:
+        """Returns the balance of 'whose' in 'ctx' at 'ctn'.
+        Only HTLCs that have been settled by that ctn are counted.
+        """
+        if ctn is None:
+            ctn = self.ctn_oldest_unrevoked(ctx_owner)
+        balance = initial_balance_msat
+        if ctn >= self.ctn_oldest_unrevoked(ctx_owner):
+            balance += self._balance_delta * whose
+            considered_sent_htlc_ids = self._maybe_active_htlc_ids[whose]
+            considered_recv_htlc_ids = self._maybe_active_htlc_ids[-whose]
+        else:  # ctn is too old; need to consider full log (slow...)
+            considered_sent_htlc_ids = self.log[whose]['settles']
+            considered_recv_htlc_ids = self.log[-whose]['settles']
+        # sent htlcs
+        for htlc_id in considered_sent_htlc_ids:
+            ctns = self.log[whose]['settles'].get(htlc_id, None)
+            if ctns is None: continue
+            if ctns[ctx_owner] is not None and ctns[ctx_owner] <= ctn:
+                htlc = self.log[whose]['adds'][htlc_id]
+                balance -= htlc.amount_msat
+        # recv htlcs
+        for htlc_id in considered_recv_htlc_ids:
+            ctns = self.log[-whose]['settles'].get(htlc_id, None)
+            if ctns is None: continue
+            if ctns[ctx_owner] is not None and ctns[ctx_owner] <= ctn:
+                htlc = self.log[-whose]['adds'][htlc_id]
+                balance += htlc.amount_msat
+        return balance
+
     def _get_htlcs_that_got_removed_exactly_at_ctn(
             self, ctn: int, *, ctx_owner: HTLCOwner, htlc_proposer: HTLCOwner, log_action: str,
     ) -> Sequence[UpdateAddHtlc]:

electrum/tests/test_lnpeer.py

@@ -306,14 +306,14 @@ class TestPeer(ElectrumTestCase):
         with self.assertRaises(concurrent.futures.CancelledError):
             run(f())
 
-    @unittest.skip("too expensive")
+    #@unittest.skip("too expensive")
     #@needs_test_with_all_chacha20_implementations
     def test_payments_stresstest(self):
         alice_channel, bob_channel = create_test_channels()
         p1, p2, w1, w2, _q1, _q2 = self.prepare_peers(alice_channel, bob_channel)
         alice_init_balance_msat = alice_channel.balance(HTLCOwner.LOCAL)
         bob_init_balance_msat = bob_channel.balance(HTLCOwner.LOCAL)
-        num_payments = 1000
+        num_payments = 50
         #pay_reqs1 = [self.prepare_invoice(w1, amount_sat=1) for i in range(num_payments)]
         pay_reqs2 = [self.prepare_invoice(w2, amount_sat=1) for i in range(num_payments)]
         max_htlcs_in_flight = asyncio.Semaphore(5)