starred/chia-blockchain
1
0
Fork 0
mirror of https://github.com/Chia-Network/chia-blockchain.git synced 2026-08-24 10:05:29 -05:00
Code Issues Packages Projects Releases Wiki Activity

CHIA-4017 Improve active requests tracking (#20748)

Browse Source 
Improve active requests tracking.
This commit is contained in:
Amine Khaldi
2026-04-02 11:13:43 -05:00
committed by GitHub
parent 04600105ae
commit fc79ba0c35
2 changed files with 197 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
chia/_tests/core/server/test_server.py
+137 -2
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import logging
from collections.abc import Callable
from dataclasses import dataclass
@@ -7,7 +8,7 @@ from typing import ClassVar, cast
import pytest
from chia_rs.sized_bytes import bytes32
from chia_rs.sized_ints import int16, uint32
from chia_rs.sized_ints import int16, uint8, uint16, uint32
from packaging.version import Version
from chia import __version__
@@ -18,7 +19,7 @@ from chia._tests.util.time_out_assert import time_out_assert
from chia.full_node.full_node_api import FullNodeAPI
from chia.full_node.start_full_node import create_full_node_service
from chia.protocols.full_node_protocol import RejectBlock, RequestBlock, RequestTransaction
from chia.protocols.outbound_message import NodeType, make_msg
from chia.protocols.outbound_message import Message, NodeType, make_msg
from chia.protocols.protocol_message_types import ProtocolMessageTypes
from chia.protocols.shared_protocol import Error, protocol_version
from chia.protocols.wallet_protocol import RejectHeaderRequest
@@ -27,6 +28,7 @@ from chia.server.server import ChiaServer
from chia.server.ssl_context import chia_ssl_ca_paths, private_ssl_ca_paths
from chia.server.ws_connection import WSChiaConnection, error_response_version, sanitize_version_string
from chia.simulator.block_tools import BlockTools
from chia.simulator.full_node_simulator import FullNodeSimulator
from chia.types.peer_info import PeerInfo
from chia.util.errors import ApiError, Err
from chia.wallet.start_wallet import create_wallet_service
@@ -305,3 +307,136 @@ async def test_connection_closed_banning(bt: BlockTools, caplog: pytest.LogCaptu
trusted_peer = FakeConnection(peer_info=PeerInfo("34.34.34.34", 8444))
await my_test_server.connection_closed(cast(WSChiaConnection, trusted_peer), ban_time=60)
assert f"Trying to ban trusted peer {trusted_peer.peer_info.host} for 60, but will not ban" in caplog.text
@pytest.mark.anyio
@pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.HARD_FORK_2_0], reason="irrelevant")
async def test_route_incoming_message(
one_node_one_block: tuple[FullNodeSimulator, ChiaServer, BlockTools], self_hostname: str
) -> None:
"""
Covers the scenarios where incoming messages are responses to active
requests, responses to timed out requests (late) and the rest.
"""
_, server, _ = one_node_one_block
wsc, _ = await add_dummy_connection_wsc(server, self_hostname, 1337)
test_response_type = uint8(ProtocolMessageTypes.respond_block.value)
# In time message gets handled via the response route
response_event = asyncio.Event()
pending_request_id = uint16(5)
wsc.pending_requests[pending_request_id] = response_event
in_time_msg = Message(type=test_response_type, id=pending_request_id, data=b"")
await wsc._route_incoming_message(in_time_msg)
assert wsc.request_results[pending_request_id] == in_time_msg
assert response_event.is_set()
# Timed out message gets dropped
timed_out_request_id = uint16(6)
wsc.timed_out_requests.add(timed_out_request_id)
timed_out_msg = Message(type=test_response_type, id=timed_out_request_id, data=b"")
incoming_queue_size = wsc.incoming_queue.qsize()
await wsc._route_incoming_message(timed_out_msg)
assert timed_out_request_id not in wsc.timed_out_requests
# The incoming queue size doesn't increase as the message got dropped
assert wsc.incoming_queue.qsize() == incoming_queue_size
# Other messages are forwarded to the incoming queue
other_msg = Message(test_response_type, uint16(7), b"")
await wsc._route_incoming_message(other_msg)
assert wsc.incoming_queue.qsize() == incoming_queue_size + 1
@pytest.mark.anyio
@pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.HARD_FORK_2_0], reason="irrelevant")
@pytest.mark.parametrize("is_outbound, range_start, range_end", [(True, 0, 2**15 - 1), (False, 2**15, 2**16 - 1)])
async def test_select_request_nonce(
one_node_one_block: tuple[FullNodeSimulator, ChiaServer, BlockTools],
self_hostname: str,
is_outbound: bool,
range_start: int,
range_end: int,
) -> None:
"""
Covers the scenarios of skipping active nonces, selecting nonces that are
available in range, wraps occurring and no available nonces. We control
inbound/outbound via `is_outbound`.
"""
_, server, _ = one_node_one_block
wsc, _ = await add_dummy_connection_wsc(server, self_hostname, 1337)
wsc.is_outbound = is_outbound
# Skip used nonces
wsc.pending_requests[uint16(range_start)] = asyncio.Event()
wsc.timed_out_requests.add(uint16(range_start + 1))
wsc.request_nonce = uint16(range_start)
assert wsc._select_request_nonce() == uint16(range_start + 2)
assert wsc.request_nonce == uint16(range_start + 3)
# Make sure that next selection skips properly if `request_nonce` points to
# a reserved nonce after successful allocation.
wsc.pending_requests.clear()
wsc.timed_out_requests.clear()
wsc.pending_requests[uint16(range_start)] = asyncio.Event()
wsc.timed_out_requests.add(uint16(range_start + 1))
wsc.pending_requests[uint16(range_start + 3)] = asyncio.Event()
wsc.request_nonce = uint16(range_start)
assert wsc._select_request_nonce() == uint16(range_start + 2)
# This points to a reserved nonce
assert wsc.request_nonce == uint16(range_start + 3)
wsc.timed_out_requests.add(uint16(range_start + 4))
# Next selection should skip the reserved nonce and the timed out nonce
assert wsc._select_request_nonce() == uint16(range_start + 5)
assert wsc.request_nonce == uint16(range_start + 6)
# Nonce in range and unused
wsc.request_nonce = uint16(range_start + 2)
wsc.pending_requests.clear()
wsc.timed_out_requests.clear()
assert wsc._select_request_nonce() == uint16(range_start + 2)
assert wsc.request_nonce == uint16(range_start + 3)
# Wrap stays within inbound/outbound range
wsc.pending_requests.clear()
wsc.timed_out_requests.clear()
wsc.pending_requests[uint16(range_end)] = asyncio.Event()
wsc.request_nonce = uint16(range_end)
assert wsc._select_request_nonce() == uint16(range_start)
assert wsc.request_nonce == uint16(range_start + 1)
# No nonces available
wsc.pending_requests.clear()
wsc.timed_out_requests.clear()
for nonce in range(range_start, range_end + 1):
wsc.pending_requests[uint16(nonce)] = asyncio.Event()
wsc.request_nonce = uint16(range_start)
assert wsc._select_request_nonce() is None
# Make sure we don't advance in this case
assert wsc.request_nonce == uint16(range_start)
# With no nonces currently available, send a request and make sure the
# connection gets closed.
msg = make_msg(ProtocolMessageTypes.request_block, RequestBlock(uint32(42), False))
response = await wsc.send_request(message_no_id=msg, timeout=1)
assert response is None
assert wsc.closed
@pytest.mark.anyio
@pytest.mark.limit_consensus_modes(allowed=[ConsensusMode.HARD_FORK_2_0], reason="irrelevant")
async def test_inbound_handler_none_msg(
one_node_one_block: tuple[FullNodeSimulator, ChiaServer, BlockTools],
self_hostname: str,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""
Covers the scenario where the inbound handler receives `None` from
`_read_one_message`.
"""
_, server, _ = one_node_one_block
wsc, peer_id = await add_dummy_connection_wsc(server, self_hostname, 1337)
await time_out_assert(5, lambda: peer_id in server.all_connections)
read_calls = 0
async def test_read_one_message() -> Message | None:
nonlocal read_calls
read_calls += 1
return None
monkeypatch.setattr(wsc, "_read_one_message", test_read_one_message)
assert wsc.inbound_task is not None
await asyncio.wait_for(wsc.inbound_task, timeout=1)
assert read_calls == 1
assert wsc.incoming_queue.qsize() == 0
await wsc.close()
chia/server/ws_connection.py
+60 -24
View File
@@ -126,6 +126,9 @@ class WSChiaConnection:
pending_requests: dict[uint16, asyncio.Event] = field(default_factory=dict, repr=False)
request_results: dict[uint16, Message] = field(default_factory=dict, repr=False)
# Nonces of `call_api` requests that timed out. Used to distinguish a late
# response from an unsolicited one.
timed_out_requests: set[uint16] = field(default_factory=set, repr=False)
closed: bool = False
connection_type: NodeType | None = None
request_nonce: uint16 = uint16(0)
@@ -530,19 +533,31 @@ class WSChiaConnection:
api_task = create_referenced_task(self._api_call(message, task_id))
self.api_tasks[task_id] = api_task
async def _route_incoming_message(self, message: Message) -> None:
if message.id in self.pending_requests:
self.request_results[message.id] = message
self.pending_requests[message.id].set()
return
if message.id in self.timed_out_requests:
# This is a late response for a timed out request. The caller is
# expected to move on.
self.timed_out_requests.discard(message.id)
self.log.info(
f"Ignoring late response {ProtocolMessageTypes(message.type).name} "
f"from {self.peer_info.host} version {self.version}"
)
return
await self.incoming_queue.put(message)
async def inbound_handler(self) -> None:
try:
while not self.closed:
message = await self._read_one_message()
if message is not None:
if message.id in self.pending_requests:
self.request_results[message.id] = message
event = self.pending_requests[message.id]
event.set()
else:
await self.incoming_queue.put(message)
else:
continue
if message is None:
# This indicates a connection shutdown so there is no point
# to continue.
break
await self._route_incoming_message(message)
except asyncio.CancelledError:
self.log.debug("Inbound_handler task cancelled")
except Exception as e:
@@ -601,6 +616,27 @@ class WSChiaConnection:
assert receive_metadata is not None, f"ApiMetadata unavailable for {recv_method}"
return receive_metadata.message_class.from_bytes(response.data)
def _next_request_nonce(self, nonce: uint16) -> uint16:
if self.is_outbound:
return uint16(nonce + 1) if nonce != uint16(2**15 - 1) else uint16(0)
return uint16(nonce + 1) if nonce != uint16(2**16 - 1) else uint16(2**15)
def _select_request_nonce(self) -> uint16 | None:
"""
Allocate a nonce that is not used by a pending or a timed out request
and advance `request_nonce` to the next candidate.
"""
start_nonce = self.request_nonce
nonce_candidate = start_nonce
while True:
if nonce_candidate not in self.pending_requests and nonce_candidate not in self.timed_out_requests:
self.request_nonce = self._next_request_nonce(nonce_candidate)
return nonce_candidate
nonce_candidate = self._next_request_nonce(nonce_candidate)
# See if we're back to the start
if nonce_candidate == start_nonce:
return None
async def send_request(self, message_no_id: Message, timeout: int) -> Message | None:
"""Sends a message and waits for a response."""
if self.closed:
@@ -611,11 +647,13 @@ class WSChiaConnection:
# The request nonce is an integer between 0 and 2**16 - 1, which is used to match requests to responses
# If is_outbound, 0 <= nonce < 2^15, else 2^15 <= nonce < 2^16
request_id = self.request_nonce
if self.is_outbound:
self.request_nonce = uint16(self.request_nonce + 1) if self.request_nonce != (2**15 - 1) else uint16(0)
else:
self.request_nonce = uint16(self.request_nonce + 1) if self.request_nonce != (2**16 - 1) else uint16(2**15)
request_id = self._select_request_nonce()
if request_id is None:
self.log.info(
f"Disconnecting peer {self.peer_info.host} version {self.version} for no available request nonces"
)
await self.close()
return None
message = Message(message_no_id.type, request_id, message_no_id.data)
assert message.id is not None
@@ -627,16 +665,14 @@ class WSChiaConnection:
except asyncio.TimeoutError:
self.log.debug(f"Request timeout: {message}")
self.pending_requests.pop(message.id)
result: Message | None = None
if message.id in self.request_results:
result = self.request_results[message.id]
assert result is not None
self.log.debug(
f"<- {ProtocolMessageTypes(result.type).name} from: {self.peer_info.host}:{self.peer_info.port}"
)
self.request_results.pop(message.id)
result = self.request_results.pop(message.id, None)
if result is None:
# This request has timed out
self.timed_out_requests.add(message.id)
self.pending_requests.pop(message.id, None)
return None
self.log.debug(f"<- {ProtocolMessageTypes(result.type).name} from: {self.peer_info.host}:{self.peer_info.port}")
self.pending_requests.pop(message.id, None)
return result
async def _wait_and_retry(self, msg: Message) -> None: