tests: test_lnpeer: test_htlc_switch_iteration_benchmark

diff --git a/tests/test_lnchannel.py b/tests/test_lnchannel.py
index 2ba4c40..6299042 100644
--- a/tests/test_lnchannel.py
+++ b/tests/test_lnchannel.py
@@ -50,7 +50,8 @@ one_bitcoin_in_msat = bitcoin.COIN * 1000
 def create_channel_state(funding_txid, funding_index, funding_sat, is_initiator,
                          local_amount, remote_amount, privkeys, other_pubkeys,
                          seed, cur, nex, other_node_id, l_dust, r_dust, l_csv,
-                         r_csv, anchor_outputs, local_max_inflight, remote_max_inflight):
+                         r_csv, anchor_outputs, local_max_inflight, remote_max_inflight,
+                         max_accepted_htlcs):
     #assert local_amount > 0
     #assert remote_amount > 0
 
@@ -71,7 +72,7 @@ def create_channel_state(funding_txid, funding_index, funding_sat, is_initiator,
                 to_self_delay=r_csv,
                 dust_limit_sat=r_dust,
                 max_htlc_value_in_flight_msat=remote_max_inflight,
-                max_accepted_htlcs=5,
+                max_accepted_htlcs=max_accepted_htlcs,
                 initial_msat=remote_amount,
                 reserve_sat=0,
                 htlc_minimum_msat=1,
@@ -91,7 +92,7 @@ def create_channel_state(funding_txid, funding_index, funding_sat, is_initiator,
                 to_self_delay=l_csv,
                 dust_limit_sat=l_dust,
                 max_htlc_value_in_flight_msat=local_max_inflight,
-                max_accepted_htlcs=5,
+                max_accepted_htlcs=max_accepted_htlcs,
                 initial_msat=local_amount,
                 reserve_sat=0,
                 per_commitment_secret_seed=seed,
@@ -133,7 +134,8 @@ def create_test_channels(*, feerate=6000, local_msat=None, remote_msat=None,
                          alice_name="alice", bob_name="bob",
                          alice_pubkey=b"\x01"*33, bob_pubkey=b"\x02"*33, random_seed=None,
                          anchor_outputs=False,
-                         local_max_inflight=None, remote_max_inflight=None):
+                         local_max_inflight=None, remote_max_inflight=None,
+                         max_accepted_htlcs=5):
     if random_seed is None:  # needed for deterministic randomness
         random_seed = os.urandom(32)
     random_gen = PRNG(random_seed)
@@ -168,7 +170,8 @@ def create_test_channels(*, feerate=6000, local_msat=None, remote_msat=None,
                 remote_amount, alice_privkeys, bob_pubkeys, alice_seed, None,
                 bob_first, other_node_id=bob_pubkey, l_dust=200, r_dust=1300,
                 l_csv=5, r_csv=4, anchor_outputs=anchor_outputs,
-                local_max_inflight=local_max_inflight, remote_max_inflight=remote_max_inflight
+                local_max_inflight=local_max_inflight, remote_max_inflight=remote_max_inflight,
+                max_accepted_htlcs=max_accepted_htlcs,
             ),
             name=f"{alice_name}->{bob_name}",
             initial_feerate=feerate),
@@ -178,7 +181,8 @@ def create_test_channels(*, feerate=6000, local_msat=None, remote_msat=None,
                 local_amount, bob_privkeys, alice_pubkeys, bob_seed, None,
                 alice_first, other_node_id=alice_pubkey, l_dust=1300, r_dust=200,
                 l_csv=4, r_csv=5, anchor_outputs=anchor_outputs,
-                local_max_inflight=remote_max_inflight, remote_max_inflight=local_max_inflight
+                local_max_inflight=remote_max_inflight, remote_max_inflight=local_max_inflight,
+                max_accepted_htlcs=max_accepted_htlcs,
             ),
             name=f"{bob_name}->{alice_name}",
             initial_feerate=feerate)
diff --git a/tests/test_lnpeer.py b/tests/test_lnpeer.py
index db38087..204ebd6 100644
--- a/tests/test_lnpeer.py
+++ b/tests/test_lnpeer.py
@@ -14,6 +14,7 @@ from unittest import mock
 from typing import Iterable, NamedTuple, Tuple, List, Dict, Sequence
 from types import MappingProxyType
 import time
+import statistics
 
 from aiorpcx import timeout_after, TaskTimeout
 from electrum_ecc import ECPrivkey
@@ -1886,6 +1887,73 @@ class TestPeerDirect(TestPeer):
         for _test_trampoline in [False, True]:
             await run_test(_test_trampoline)
 
+    async def test_htlc_switch_iteration_benchmark(self):
+        """Test how long a call to _run_htlc_switch_iteration takes with 10 trampoline
+        mpp sets of 1 htlc each. Raise if it takes longer than 20ms (median).
+        To create flamegraph with py-spy raise NUM_ITERATIONS to 1000 (for more samples) then run:
+        $ py-spy record -o flamegraph.svg --subprocesses -- python -m pytest tests/test_lnpeer.py::TestPeerDirect::test_htlc_switch_iteration_benchmark
+        """
+        NUM_ITERATIONS = 25
+        alice_channel, bob_channel = create_test_channels(max_accepted_htlcs=20)
+        alice_p, bob_p, alice_w, bob_w, _q1, _q2 = self.prepare_peers(alice_channel, bob_channel)
+
+        await self._activate_trampoline(alice_w)
+        electrum.trampoline._TRAMPOLINE_NODES_UNITTESTS = {
+            'bob': LNPeerAddr(host="127.0.0.1", port=9735, pubkey=bob_w.node_keypair.pubkey),
+        }
+
+        # create 10 invoices (10 pending htlc sets with 1 htlc each)
+        invoices = []  # type: list[tuple[LnAddr, Invoice]]
+        for i in range(10):
+            lnaddr, pay_req = self.prepare_invoice(bob_w)
+            # prevent bob from settling so that htlc switch will have to iterate through all pending htlcs
+            bob_w.dont_settle_htlcs[pay_req.rhash] = None
+            invoices.append((lnaddr, pay_req))
+        self.assertEqual(len(invoices), 10, msg=len(invoices))
+
+        iterations = []
+        do_benchmark = False
+        _run_bob_htlc_switch_iteration = bob_p._run_htlc_switch_iteration
+        def timed_htlc_switch_iteration():
+            start = time.perf_counter()
+            _run_bob_htlc_switch_iteration()
+            duration = time.perf_counter() - start
+            if do_benchmark:
+                iterations.append(duration)
+        bob_p._run_htlc_switch_iteration = timed_htlc_switch_iteration
+
+        async def benchmark_htlc_switch_iterations():
+            waited = 0
+            while not len(bob_w.received_mpp_htlcs) == 10 :
+                waited += 0.1
+                await asyncio.sleep(0.1)
+                if waited > 2:
+                    raise TimeoutError()
+            nonlocal do_benchmark
+            do_benchmark = True
+            while len(iterations) < NUM_ITERATIONS:
+                await asyncio.sleep(0.05)
+            # average = sum(iterations) / len(iterations)
+            median_duration = statistics.median(iterations)
+            res = f"median duration per htlc switch iteration: {median_duration:.6f}s over {len(iterations)=}"
+            self.logger.info(res)
+            self.assertLess(median_duration, 0.02, msg=res)
+            raise SuccessfulTest()
+
+        async def f():
+            async with OldTaskGroup() as group:
+                await group.spawn(alice_p._message_loop())
+                await group.spawn(alice_p.htlc_switch())
+                await group.spawn(bob_p._message_loop())
+                await group.spawn(bob_p.htlc_switch())
+                await asyncio.sleep(0.01)
+                for _lnaddr, req in invoices:
+                    await group.spawn(alice_w.pay_invoice(req))
+                await benchmark_htlc_switch_iterations()
+
+        with self.assertRaises(SuccessfulTest):
+            await f()
+
 
 class TestPeerForwarding(TestPeer):