CHIA-4261: Use bisect for request_ses_hashes SES lookup (#21246)

chia/_tests/core/full_node/test_request_ses_hashes.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import pytest
+from chia_rs.sized_ints import uint32
+
+from chia.full_node.full_node_api import ses_intervals_for_range
+
+
+@pytest.mark.parametrize(
+    "ses_heights, start, end, expected",
+    [
+        # Too few SES entries → empty.
+        ([], 10, 20, []),
+        ([100], 50, 150, []),
+        # start before the first SES height.
+        ([100, 200, 300, 400], 50, 150, []),
+        # start after the last SES height.
+        ([100, 200, 300, 400], 400, 500, []),
+        ([100, 200, 300, 400], 500, 600, []),
+        # Entire request inside one SES interval.
+        ([100, 200, 300, 400], 100, 150, [(100, 200)]),
+        ([100, 200, 300, 400], 150, 199, [(100, 200)]),
+        # end on the next boundary is treated as spanning (strict upper bound).
+        ([100, 200, 300, 400], 150, 200, [(100, 200), (200, 300)]),
+        # Request spans two SES intervals.
+        ([100, 200, 300, 400], 150, 250, [(100, 200), (200, 300)]),
+        # start exactly on an SES boundary.
+        ([100, 200, 300, 400], 200, 250, [(200, 300)]),
+        ([100, 200, 300, 400], 200, 350, [(200, 300), (300, 400)]),
+        # Last interval: cannot append a following SES even if end is past it.
+        ([100, 200, 300, 400], 350, 999, [(300, 400)]),
+    ],
+    ids=[
+        "empty_heights",
+        "single_height",
+        "before_first",
+        "at_last_height",
+        "after_last",
+        "first_interval_at_start",
+        "first_interval_interior",
+        "end_on_next_boundary_spans",
+        "spans_two",
+        "on_boundary_same_interval",
+        "on_boundary_spans",
+        "last_interval_only",
+    ],
+)
+def test_ses_intervals_for_range(
+    ses_heights: list[int],
+    start: int,
+    end: int,
+    expected: list[tuple[int, int]],
+) -> None:
+    result = ses_intervals_for_range(
+        [uint32(h) for h in ses_heights],
+        uint32(start),
+        uint32(end),
+    )
+    assert [(int(a), int(b)) for a, b in result] == expected

chia/full_node/full_node_api.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import asyncio
+import bisect
 import logging
 import time
 import traceback
-from collections.abc import Collection
+from collections.abc import Collection, Sequence
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, ClassVar, cast
 
@@ -24,7 +25,6 @@ from chia_rs import (
     PoolTarget,
     RespondToPhUpdates,
     RewardChainBlockUnfinished,
-    SubEpochSummary,
     UnfinishedBlock,
     additions_and_removals,
     get_flags_for_height_and_constants,
@@ -87,6 +87,36 @@ MAX_COIN_HASHES_PER_REQUEST = 50
 MAX_COINS_MAP_SIZE = 100
 
 
+def ses_intervals_for_range(
+    ses_heights: Sequence[uint32],
+    start_height: uint32,
+    end_height: uint32,
+) -> list[tuple[uint32, uint32]]:
+    """
+    Return 0-2 SES height intervals (start, next) covering [start_height, end_height].
+
+    Uses bisect instead of a linear scan (mainnet has tens of thousands of SES heights).
+    """
+    if len(ses_heights) < 2:
+        return []
+
+    idx = bisect.bisect_right(ses_heights, start_height) - 1
+    if not (0 <= idx < len(ses_heights) - 1):
+        return []
+
+    ses_start_height = ses_heights[idx]
+    next_ses_height = ses_heights[idx + 1]
+    if not (ses_start_height <= start_height < next_ses_height):
+        return []
+
+    intervals: list[tuple[uint32, uint32]] = [(ses_start_height, next_ses_height)]
+    if not (ses_start_height < end_height < next_ses_height) and idx < len(ses_heights) - 2:
+        # Request spans two SES intervals.
+        next_next_height = ses_heights[idx + 2]
+        intervals.append((next_ses_height, next_next_height))
+    return intervals
+
+
 async def tx_request_and_timeout(full_node: FullNode, transaction_id: bytes32, task_id: bytes32) -> None:
     """
     Request a transaction from peers that advertised it, until we either
@@ -1947,32 +1977,9 @@ class FullNodeAPI:
         """Returns the start and end height of a sub-epoch for the height specified in request"""
 
         ses_height = self.full_node.blockchain.get_ses_heights()
-        start_height = request.start_height
-        end_height = request.end_height
-        ses_hash_heights = []
-        ses_reward_hashes = []
-
-        for idx, ses_start_height in enumerate(ses_height):
-            if idx == len(ses_height) - 1:
-                break
-
-            next_ses_height = ses_height[idx + 1]
-            # start_ses_hash
-            if ses_start_height <= start_height < next_ses_height:
-                ses_hash_heights.append([ses_start_height, next_ses_height])
-                ses: SubEpochSummary = self.full_node.blockchain.get_ses(ses_start_height)
-                ses_reward_hashes.append(ses.reward_chain_hash)
-                if ses_start_height < end_height < next_ses_height:
-                    break
-                else:
-                    if idx == len(ses_height) - 2:
-                        break
-                    # else add extra ses as request start <-> end spans two ses
-                    next_next_height = ses_height[idx + 2]
-                    ses_hash_heights.append([next_ses_height, next_next_height])
-                    nex_ses: SubEpochSummary = self.full_node.blockchain.get_ses(next_ses_height)
-                    ses_reward_hashes.append(nex_ses.reward_chain_hash)
-                    break
+        intervals = ses_intervals_for_range(ses_height, request.start_height, request.end_height)
+        ses_hash_heights = [[start, end] for start, end in intervals]
+        ses_reward_hashes = [self.full_node.blockchain.get_ses(start).reward_chain_hash for start, _end in intervals]
 
         response = RespondSESInfo(ses_reward_hashes, ses_hash_heights)
         msg = make_msg(ProtocolMessageTypes.respond_ses_hashes, response)