# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for the unified bidirectional P2PSession. A P2PSession owns a single ControlConnection and dispatches every message type — both client-role (FetchMsg / TransferDoneMsg / AbortAck) and server-role (FetchMsg / TransferDoneMsg / AbortAck from the peer's perspective). These tests exercise both flows independently and the bidirectional case where one session simultaneously serves a fetch and completes its own load. """ from __future__ import annotations import time from collections.abc import Sequence import numpy as np import pytest from vllm.v1.kv_offload.base import ( LookupResult, OffloadKey, ReqContext, RequestOffloadingContext, ) from vllm.v1.kv_offload.tiering.base import TransferJob from vllm.v1.kv_offload.tiering.p2p.session import ( LoadResult, P2PSession, StoreResult, ) from vllm.v1.kv_offload.tiering.p2p.session.client import ( _ABORT_ACK_TIMEOUT_S, _LOAD_TIMEOUT_S, ) from vllm.v1.kv_offload.tiering.p2p.session.protocol import ( TYPE_KEY, AbortAckMsg, AbortFetchMsg, ConnectAckMsg, ConnectMsg, DisconnectMsg, FetchMsg, LookupMsg, LookupRespMsg, TransferDoneMsg, ) from vllm.v1.kv_offload.tiering.p2p.session.server import ( _CANCEL_DRAIN_TIMEOUT_S, _InflightXfer, _OutboundRequestState, ) from vllm.v1.kv_offload.tiering.p2p.session.session import ( _MAX_CONSECUTIVE_DISPATCH_ERRORS, ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- # Shared PYTHONHASHSEED used by the session under test and the fake peer's # ConnectMsg so the handshake succeeds unless a test overrides one side. _DEFAULT_HASH_SEED = "0" class FakeDataTransport: """Minimal fake DataTransport for testing sessions.""" def __init__( self, base_addr: int = 0x1000, num_blocks: int = 16, block_len: int = 4096, config_fingerprint: str = "", ) -> None: self._base_addr = base_addr self._num_blocks = num_blocks self._block_len = block_len self._config_fingerprint = config_fingerprint self._remote_peers: dict[str, dict] = {} self._transfers: dict[int, tuple] = {} self._next_id = 0 self._poll_done: list[int] = [] self._poll_failed: list[int] = [] self._cancel_still_inflight: set[int] = set() self._cancel_calls: list[tuple[list[int], str]] = [] @property def base_addr(self) -> int: return self._base_addr @property def num_blocks(self) -> int: return self._num_blocks @property def block_len(self) -> int: return self._block_len @property def config_fingerprint(self) -> str: return self._config_fingerprint def get_agent_metadata(self) -> bytes: return b"fake-metadata" def add_remote_peer( self, peer_id, agent_metadata, base_addr, num_blocks, block_len ) -> None: self._remote_peers[peer_id] = { ConnectMsg.AGENT_METADATA: agent_metadata, ConnectMsg.BASE_ADDR: base_addr, "num_blocks": num_blocks, "block_len": block_len, } def remove_remote_peer(self, peer_id: str) -> None: self._remote_peers.pop(peer_id, None) def write_blocks(self, peer_id, local_idxs, remote_idxs) -> int | None: if peer_id not in self._remote_peers: return None tid = self._next_id self._next_id += 1 self._transfers[tid] = (peer_id, local_idxs, remote_idxs) return tid def poll(self, peer_id=None): from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult result = PollResult(done=list(self._poll_done), failed=list(self._poll_failed)) self._poll_done.clear() self._poll_failed.clear() return result def cancel(self, transfer_ids, mode: str = "immediate") -> list[int]: ids = list(transfer_ids) self._cancel_calls.append((ids, mode)) if mode == "wait": still: list[int] = [] for tid in ids: if tid in self._cancel_still_inflight: still.append(tid) else: self._transfers.pop(tid, None) return still for tid in ids: self._transfers.pop(tid, None) self._cancel_still_inflight.discard(tid) return [] def close(self) -> None: pass class FakeConnection: """Fake ControlConnection that captures sent messages.""" def __init__(self, peer_id: str = "peer:8000") -> None: self.peer_id = peer_id self._inbox: list[dict] = [] self._sent: list[dict] = [] self._closed = False # When True, send() raises to simulate a broken/dead connection. self.fail_send = False @property def alive(self) -> bool: return not self._closed def send(self, msg: dict) -> None: if self.fail_send: raise ConnectionError("simulated dead connection") self._sent.append(msg) def recv(self) -> list[dict]: msgs = self._inbox self._inbox = [] return msgs def enqueue(self, msg: dict) -> None: self._inbox.append(msg) def mark_dead(self) -> None: self._closed = True def close(self) -> None: self._closed = True def _peer_connect_msg( peer_id: str = "peer:8000", block_len: int = 4096, fingerprint: str | None = None, hash_seed: str = _DEFAULT_HASH_SEED, ) -> dict: """Build a ConnectMsg as if the peer sent it.""" msg = { TYPE_KEY: ConnectMsg.TYPE, ConnectMsg.PEER_ID: peer_id, ConnectMsg.AGENT_METADATA: b"peer-metadata", ConnectMsg.BASE_ADDR: 0x2000, ConnectMsg.NUM_BLOCKS: 16, ConnectMsg.BLOCK_LEN: block_len, ConnectMsg.HASH_SEED: hash_seed, } if fingerprint is not None: msg[ConnectMsg.CONFIG_FINGERPRINT] = fingerprint return msg class FakeParent: """Configurable :class:`ParentManager` for server-role tests. ``stored`` is the dict of ready blocks (key → primary block_id). ``pending`` and ``retry`` script the first lookup() result for those keys; subsequent lookups behave normally (a key that promised HIT_PENDING / RETRY can later be promoted to HIT by adding it to ``stored`` and removing it from ``pending``/``retry``). ``calls`` captures every parent invocation in order for assertions. Injected per-step via ``session.serve_external_requests(parent)`` — not held by the session, matching how ``TieringOffloadingManager`` hands the tier a handle valid only for that call. """ def __init__( self, stored: dict[OffloadKey, int] | None = None, pending: set[OffloadKey] | None = None, retry: set[OffloadKey] | None = None, ) -> None: self.stored: dict[OffloadKey, int] = dict(stored or {}) self.pending: set[OffloadKey] = set(pending or ()) self.retry: set[OffloadKey] = set(retry or ()) self._next_job_id: int = 1000 self.calls: list[tuple] = [] def on_new_request(self, ctx: ReqContext) -> RequestOffloadingContext: self.calls.append(("on_new_request", ctx.req_id)) return RequestOffloadingContext() def lookup(self, key: OffloadKey, ctx: ReqContext) -> LookupResult: self.calls.append(("lookup", key, ctx.req_id)) if key in self.pending: return LookupResult.HIT_PENDING if key in self.retry: return LookupResult.RETRY if key in self.stored: return LookupResult.HIT return LookupResult.MISS def create_store_job( self, keys: Sequence[OffloadKey], ctx: ReqContext, ) -> TransferJob: keys_list = list(keys) self.calls.append(("create_store_job", tuple(keys_list), ctx.req_id)) block_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32) job_id = self._next_job_id self._next_job_id += 1 return TransferJob( job_id=job_id, keys=keys_list, block_ids=block_ids, is_promotion=False, req_context=ctx, ) def on_request_finished(self, ctx: ReqContext) -> None: self.calls.append(("on_request_finished", ctx.req_id)) def _make_session( conn: FakeConnection | None = None, transport: FakeDataTransport | None = None, peer_id: str = "peer:8000", local_id: str = "local:9000", local_hash_seed: str = _DEFAULT_HASH_SEED, ) -> tuple[P2PSession, FakeConnection, FakeDataTransport]: if conn is None: conn = FakeConnection(peer_id=peer_id) if transport is None: transport = FakeDataTransport() session = P2PSession( peer_id=peer_id, local_id=local_id, transport=transport, # type: ignore[arg-type] local_block_len=transport.block_len, local_hash_seed=local_hash_seed, conn=conn, # type: ignore[arg-type] ) return session, conn, transport def _serve(session: P2PSession, parent: FakeParent) -> None: """Resolve enqueued inbound lookups, as the manager does each step.""" session.serve_external_requests(parent) # type: ignore[arg-type] def _activate( session: P2PSession, conn: FakeConnection, peer_id: str = "peer:8000" ) -> None: """Drive the handshake: peer sends ConnectMsg + ConnectAckMsg.""" conn.enqueue(_peer_connect_msg(peer_id=peer_id)) conn.enqueue({TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: peer_id}) session.poll() # --- Accessors for the server role's consolidated per-request state. # ServerRole keeps one _ServerRequestState per kv_request_id; entries are # garbage-collected once fully idle, so "no outbound"/"no abort" reads as # either a missing entry or a None field. These helpers paper over that. def _client_load(session: P2PSession, kv_request_id: str): """The single in-flight load of a kv_request_id (loads are per-round).""" loads = session._client._requests[kv_request_id].loads assert len(loads) == 1 return next(iter(loads.values())) def _srv_outbound(session: P2PSession, kv_request_id: str): """Serve-side round for a kv_request_id, or None (idle / GC'd). Rounds are keyed by wire round_seq; surfaces the demanded round when a fetch has bound one, else any parked supply round. """ st = session._server._requests.get(kv_request_id) if st is None or not st.outbound: return None for rnd in st.outbound.values(): if rnd.demand_received: return rnd return next(iter(st.outbound.values())) def _srv_lookups(session: P2PSession) -> list: """Every parked inbound _ActiveLookup across all requests.""" return [ lu for st in session._server._requests.values() for lu in st.lookups.values() ] def _srv_abort_started(session: P2PSession, kv_request_id: str) -> float | None: """Pending-abort start time for a kv_request_id, or None.""" for (kv, _), started in session._server._pending_aborts.items(): if kv == kv_request_id: return started return None def _srv_inflight_count(session: P2PSession, kv_request_id: str) -> int: """Inflight-transfer count tracked for a kv_request_id (0 if idle).""" st = session._server._requests.get(kv_request_id) return len(st.inflight_tids) if st is not None else 0 def _srv_total_inflight(session: P2PSession) -> int: """Sum of per-request inflight counts across all requests.""" return sum(len(st.inflight_tids) for st in session._server._requests.values()) # --------------------------------------------------------------------------- # Connect / handshake # --------------------------------------------------------------------------- class TestConnectHandshake: def test_connect_msg_sent_on_creation(self): """Session sends its own ConnectMsg on connection.""" session, conn, _ = _make_session() assert len(conn._sent) == 1 msg = conn._sent[0] assert msg[TYPE_KEY] == ConnectMsg.TYPE assert msg[ConnectMsg.PEER_ID] == "local:9000" assert msg[ConnectMsg.NUM_BLOCKS] == 16 assert msg[ConnectMsg.BLOCK_LEN] == 4096 assert ConnectMsg.AGENT_METADATA in msg def test_peer_connect_triggers_add_remote_and_ack(self): """Receiving ConnectMsg registers peer and replies with ConnectAck.""" session, conn, transport = _make_session() conn.enqueue(_peer_connect_msg()) session.poll() assert "peer:8000" in transport._remote_peers ack = next(m for m in conn._sent if m[TYPE_KEY] == ConnectAckMsg.TYPE) assert ack[ConnectAckMsg.PEER_ID] == "local:9000" def test_connect_ack_makes_session_ready(self): """Session.ready becomes True after ConnectAckMsg.""" session, conn, _ = _make_session() assert not session.ready conn.enqueue({TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: "peer:8000"}) session.poll() assert session.ready def test_messages_queued_before_ack_flush_on_ack(self): """Outgoing messages sent before ConnectAck are flushed after.""" session, conn, _ = _make_session() session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] ) # Before ack: only our ConnectMsg was sent. assert len(conn._sent) == 1 assert conn._sent[0][TYPE_KEY] == ConnectMsg.TYPE # Ack arrives. conn.enqueue({TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: "peer:8000"}) session.poll() # Queued fetch is now sent. assert any(m[TYPE_KEY] == FetchMsg.TYPE for m in conn._sent) def test_block_len_mismatch_marks_dead(self): """Mismatched block_len rejects peer and marks connection dead.""" session, conn, transport = _make_session() conn.enqueue(_peer_connect_msg(block_len=8192)) # mismatch session.poll() assert "peer:8000" not in transport._remote_peers assert not session.alive def test_config_fingerprint_mismatch_marks_dead(self): """Mismatched config fingerprint rejects peer.""" transport = FakeDataTransport(config_fingerprint="abc123") session, conn, _ = _make_session(transport=transport) conn.enqueue(_peer_connect_msg(fingerprint="different")) session.poll() assert "peer:8000" not in transport._remote_peers assert not session.alive def test_config_fingerprint_match_succeeds(self): """Matching fingerprints register the peer.""" transport = FakeDataTransport(config_fingerprint="same_fp") session, conn, _ = _make_session(transport=transport) conn.enqueue(_peer_connect_msg(fingerprint="same_fp")) session.poll() assert "peer:8000" in transport._remote_peers assert session.alive def test_missing_fingerprint_allowed(self): """Missing fingerprint on either side is allowed.""" transport = FakeDataTransport(config_fingerprint="abc123") session, conn, _ = _make_session(transport=transport) conn.enqueue(_peer_connect_msg()) # no fingerprint session.poll() assert "peer:8000" in transport._remote_peers def test_hash_seed_mismatch_marks_dead(self): """Mismatched PYTHONHASHSEED rejects peer and marks connection dead.""" session, conn, transport = _make_session(local_hash_seed="0") conn.enqueue(_peer_connect_msg(hash_seed="12345")) # mismatch session.poll() assert "peer:8000" not in transport._remote_peers assert not session.alive assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) def test_hash_seed_match_succeeds(self): """Matching PYTHONHASHSEED registers the peer and acks.""" session, conn, transport = _make_session(local_hash_seed="12345") conn.enqueue(_peer_connect_msg(hash_seed="12345")) session.poll() assert "peer:8000" in transport._remote_peers assert session.alive assert any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) def test_hash_seed_advertised_in_connect_msg(self): """Session advertises its own PYTHONHASHSEED in the ConnectMsg.""" _, conn, _ = _make_session(local_hash_seed="777") assert conn._sent[0][ConnectMsg.HASH_SEED] == "777" # --------------------------------------------------------------------------- # Client-role flows # --------------------------------------------------------------------------- class TestClientFlows: def test_request_blocks_sends_fetch(self): session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k1", b"k2"], block_ids=[0, 1] ) lookup = conn._sent[-1] assert lookup[TYPE_KEY] == FetchMsg.TYPE assert lookup[FetchMsg.KV_REQUEST_ID] == "req-1" assert lookup[FetchMsg.KEYS] == [b"k1", b"k2"] assert lookup[FetchMsg.BLOCK_INDEXES] == [0, 1] def test_transfer_done_success(self): session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] ) conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: True, } ) loads = session.poll().loads assert loads == [LoadResult(job_id=1, kv_request_id="req-1", success=True)] def test_transfer_done_failure(self): session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] ) conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: False, } ) loads = session.poll().loads assert loads == [LoadResult(job_id=1, kv_request_id="req-1", success=False)] def test_finish_request_sends_abort(self): session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] ) session.finish_request("req-1") abort = conn._sent[-1] assert abort[TYPE_KEY] == AbortFetchMsg.TYPE assert abort[AbortFetchMsg.KV_REQUEST_ID] == "req-1" def test_active_loads_work_list_tracks_in_flight(self): """collect_results / has_active_loads use the _active_loads work-list, armed when a fetch is issued and discarded exactly when its load clears — a probe-only request never enters it, and completion empties it while the entry may briefly linger for GC.""" session, conn, _ = _make_session() _activate(session, conn) client = session._client # A probe-only request has no in-flight load: not in _active_loads. session.register_lookup("req-probe", b"hp") assert client._active_loads == set() assert client.has_active_loads is False # Issuing a fetch arms the work-list. session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] ) assert client._active_loads == {"req-1"} assert client.has_active_loads is True # Completion clears the load and discards it from the work-list. conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: True, } ) session.poll() assert client._active_loads == set() assert client.has_active_loads is False def test_load_timeout_sends_abort(self): session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"k"], block_ids=[0] ) _client_load(session, "req-1").submitted_at = time.monotonic() - 60.0 session.poll() abort = conn._sent[-1] assert abort[TYPE_KEY] == AbortFetchMsg.TYPE def test_load_abort_ack_timeout_surfaces_failure(self): """After load timeout sends AbortFetch, if no AbortAck arrives within _ABORT_ACK_TIMEOUT_S the request is surfaced as failed and removed from _requests — the engine cannot wait forever on a peer that won't ack. """ session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=7, kv_request_id="req-7", keys=[b"k"], block_ids=[0] ) # 1) Trip the load timeout to send AbortFetch and stamp aborted_at. _client_load(session, "req-7").submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) loads = session.poll().loads assert loads == [] assert any( m.get(TYPE_KEY) == AbortFetchMsg.TYPE and m[AbortFetchMsg.KV_REQUEST_ID] == "req-7" for m in conn._sent ) assert _client_load(session, "req-7").aborted_at is not None # 2) Now backdate aborted_at past the abort-ack timeout. No ack ever # arrived from the peer. _client_load(session, "req-7").aborted_at = ( time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0 ) loads = session.poll().loads assert loads == [LoadResult(job_id=7, kv_request_id="req-7", success=False)] assert "req-7" not in session._client._requests def test_load_abort_ack_clears_request(self): """After load timeout sends AbortFetch, an arriving AbortAckMsg from the peer surfaces the failure cleanly and removes the request from _requests — covers the on_abort_ack arrival path.""" session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=8, kv_request_id="req-8", keys=[b"k"], block_ids=[0] ) _client_load(session, "req-8").submitted_at = ( time.monotonic() - _LOAD_TIMEOUT_S - 1.0 ) # First poll: AbortFetch goes out. session.poll() assert _client_load(session, "req-8").aborted_at is not None # Peer acks the abort. conn.enqueue( { TYPE_KEY: AbortAckMsg.TYPE, AbortAckMsg.ROUND_SEQ: 0, AbortAckMsg.KV_REQUEST_ID: "req-8", } ) loads = session.poll().loads assert loads == [LoadResult(job_id=8, kv_request_id="req-8", success=False)] assert "req-8" not in session._client._requests # --------------------------------------------------------------------------- # Symmetric-P2P lookup flow (do_p2p_fetch) # --------------------------------------------------------------------------- class TestLookupFlow: """Consumer-side state machine for do_p2p_fetch lookups.""" def test_aggregate_flush_resolve_round_trip(self): """register_lookup → flush sends one LookupMsg → response resolves entries → register_lookup returns the cached bool on every call, and repeat probes never re-issue a LookupMsg.""" session, conn, _ = _make_session() _activate(session, conn) # Aggregate two keys for the same kv_request_id; both return None. assert session.register_lookup("req-1", b"hA") is None assert session.register_lookup("req-1", b"hB") is None # Flush sends one LookupMsg with both keys. sent_before = len(conn._sent) session.flush_pending_lookups() new = conn._sent[sent_before:] assert len(new) == 1 msg = new[0] assert msg[TYPE_KEY] == LookupMsg.TYPE assert msg[LookupMsg.KV_REQUEST_ID] == "req-1" assert sorted(msg[LookupMsg.KEYS]) == [b"hA", b"hB"] # Idempotent re-flush sends nothing — the entries are now in-flight. sent_before = len(conn._sent) session.flush_pending_lookups() assert conn._sent[sent_before:] == [] # While in-flight, register_lookup keeps returning None. assert session.register_lookup("req-1", b"hA") is None # Peer answers: hA hit, hB miss. conn.enqueue( { TYPE_KEY: LookupRespMsg.TYPE, LookupRespMsg.KV_REQUEST_ID: "req-1", LookupRespMsg.KEYS: [b"hA", b"hB"], LookupRespMsg.HITS: [True, False], } ) session.poll() # register_lookup returns the resolved bool. assert session.register_lookup("req-1", b"hA") is True assert session.register_lookup("req-1", b"hB") is False # The entry is cached, not popped: repeat probes keep returning the # same result and never re-queue the key, so a flush sends nothing. assert session.register_lookup("req-1", b"hA") is True assert session.register_lookup("req-1", b"hB") is False sent_before = len(conn._sent) session.flush_pending_lookups() assert [ m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE ] == [] def test_request_blocks_clears_probe_cache(self): """A resolved HIT probe is popped when its fetch is issued, so a re-scheduled request re-probes instead of trusting the stale True (the served block is unpinned and may have been evicted).""" session, conn, _ = _make_session() _activate(session, conn) # Probe hA, flush, and let the peer resolve it to a HIT. assert session.register_lookup("req-1", b"hA") is None session.flush_pending_lookups() conn.enqueue( { TYPE_KEY: LookupRespMsg.TYPE, LookupRespMsg.KV_REQUEST_ID: "req-1", LookupRespMsg.KEYS: [b"hA"], LookupRespMsg.HITS: [True], } ) session.poll() assert session.register_lookup("req-1", b"hA") is True # Fetch consumes the probe. session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"hA"], block_ids=[0] ) assert b"hA" not in session._client._requests["req-1"].probes # Re-scheduled probe of the same key is treated as brand-new: it # returns None and re-queues, so a flush emits a fresh LookupMsg. assert session.register_lookup("req-1", b"hA") is None sent_before = len(conn._sent) session.flush_pending_lookups() fresh = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] assert len(fresh) == 1 assert fresh[0][LookupMsg.KEYS] == [b"hA"] def test_flush_uses_work_list_not_full_scan(self): """flush drains a work-list rather than scanning every live request. A request with no newly-registered keys is not revisited: after a flush the work-list is empty, an idle re-flush sends nothing, and a subsequent register re-arms exactly the one affected id — even while an unrelated request stays live in ``_requests``. """ session, conn, _ = _make_session() _activate(session, conn) client = session._client # Two requests register keys; both are queued for flush. session.register_lookup("req-A", b"hA") session.register_lookup("req-B", b"hB") assert client._flush_pending == {"req-A", "req-B"} # Flush drains the work-list even though both requests stay live. session.flush_pending_lookups() assert client._flush_pending == set() assert set(client._requests) == {"req-A", "req-B"} # An idle re-flush visits nothing and sends no LookupMsg. sent_before = len(conn._sent) session.flush_pending_lookups() assert [ m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE ] == [] # A new register re-arms only that id. session.register_lookup("req-A", b"hA2") assert client._flush_pending == {"req-A"} def test_separate_lookup_msg_per_kv_request_id(self): """Hashes for different kv_request_ids flush as separate LookupMsgs.""" session, conn, _ = _make_session() _activate(session, conn) session.register_lookup("req-A", b"h1") session.register_lookup("req-B", b"h2") session.register_lookup("req-A", b"h3") sent_before = len(conn._sent) session.flush_pending_lookups() sent = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] assert len(sent) == 2 by_req = {m[LookupMsg.KV_REQUEST_ID]: m[LookupMsg.KEYS] for m in sent} assert sorted(by_req["req-A"]) == [b"h1", b"h3"] assert by_req["req-B"] == [b"h2"] def test_multiple_lookup_msgs_across_steps(self): """A request's block set may be discovered across scheduler steps: each step that registers new keys flushes its own LookupMsg for the same kv_request_id, carrying only the newly-probed keys.""" session, conn, _ = _make_session() _activate(session, conn) # Step 1: probe hA, hB. session.register_lookup("req-1", b"hA") session.register_lookup("req-1", b"hB") sent_before = len(conn._sent) session.flush_pending_lookups() first = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] assert len(first) == 1 assert first[0][LookupMsg.KV_REQUEST_ID] == "req-1" assert sorted(first[0][LookupMsg.KEYS]) == [b"hA", b"hB"] # Step 2: a new key is discovered for the same request. The # in-flight keys from step 1 are not re-sent; a second LookupMsg # goes out carrying only the newly-probed key. assert session.register_lookup("req-1", b"hA") is None # in-flight no-op session.register_lookup("req-1", b"hC") sent_before = len(conn._sent) session.flush_pending_lookups() second = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupMsg.TYPE] assert len(second) == 1 assert second[0][LookupMsg.KV_REQUEST_ID] == "req-1" assert second[0][LookupMsg.KEYS] == [b"hC"] def test_split_response_resolves_across_messages(self): """Producer may answer one LookupMsg's keys across multiple LookupRespMsgs — pairs are self-describing so each lands.""" session, conn, _ = _make_session() _activate(session, conn) session.register_lookup("req-1", b"hA") session.register_lookup("req-1", b"hB") session.flush_pending_lookups() # Two responses, each carrying one of the two keys. conn.enqueue( { TYPE_KEY: LookupRespMsg.TYPE, LookupRespMsg.KV_REQUEST_ID: "req-1", LookupRespMsg.KEYS: [b"hA"], LookupRespMsg.HITS: [True], } ) conn.enqueue( { TYPE_KEY: LookupRespMsg.TYPE, LookupRespMsg.KV_REQUEST_ID: "req-1", LookupRespMsg.KEYS: [b"hB"], LookupRespMsg.HITS: [False], } ) session.poll() assert session.register_lookup("req-1", b"hA") is True assert session.register_lookup("req-1", b"hB") is False def test_finish_request_cancels_pending_lookups(self): """finish_request drops every pending lookup for the kv_request_id.""" session, conn, _ = _make_session() _activate(session, conn) session.register_lookup("req-1", b"hA") session.register_lookup("req-1", b"hB") session.register_lookup("req-2", b"hC") session.finish_request("req-1") # req-1 entries gone, req-2 untouched. assert "req-1" not in session._client._requests assert b"hC" in session._client._requests["req-2"].probes def test_finish_after_flushed_lookup_sends_empty_fetch(self): """LookupMsg flushed but no FetchMsg sent (all-miss case) → finish_request emits an empty FetchMsg so the peer can drop its lookup state and call parent.on_request_finished.""" session, conn, _ = _make_session() _activate(session, conn) session.register_lookup("req-1", b"hA") session.flush_pending_lookups() sent_before = len(conn._sent) session.finish_request("req-1") fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] assert len(fetches) == 1 assert fetches[0][FetchMsg.KV_REQUEST_ID] == "req-1" assert fetches[0][FetchMsg.KEYS] == [] assert fetches[0][FetchMsg.BLOCK_INDEXES] == [] def test_finish_without_flushed_lookup_sends_no_fetch(self): """No LookupMsg was ever sent → finish_request must not emit an empty FetchMsg (the peer has no state to release).""" session, conn, _ = _make_session() _activate(session, conn) # Register but never flush. session.register_lookup("req-1", b"hA") sent_before = len(conn._sent) session.finish_request("req-1") fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] assert fetches == [] def test_finish_after_real_fetch_sends_no_second_fetch(self): """A real FetchMsg was already sent for the id → finish_request must not emit a second (empty) FetchMsg.""" session, conn, _ = _make_session() _activate(session, conn) session.register_lookup("req-1", b"hA") session.flush_pending_lookups() # Resolve the probe to a HIT before fetching, as the manager only # loads confirmed hits (an unresolved probe yields RETRY). conn.enqueue( { TYPE_KEY: LookupRespMsg.TYPE, LookupRespMsg.KV_REQUEST_ID: "req-1", LookupRespMsg.KEYS: [b"hA"], LookupRespMsg.HITS: [True], } ) session.poll() session.request_blocks( job_id=1, kv_request_id="req-1", keys=[b"hA"], block_ids=[7] ) sent_before = len(conn._sent) session.finish_request("req-1") fetches = [m for m in conn._sent[sent_before:] if m[TYPE_KEY] == FetchMsg.TYPE] assert fetches == [] def test_server_lookup_deferred_until_serve_then_all_misses(self): """``poll()`` only enqueues an inbound LookupMsg — no response is sent until ``serve_external_requests``. With an all-miss parent the aggregated LookupRespMsg carries the same keys and ``hits=[False, ...]``.""" session, conn, _ = _make_session() _activate(session, conn) sent_before = len(conn._sent) conn.enqueue( { TYPE_KEY: LookupMsg.TYPE, LookupMsg.ROUND_SEQ: 0, LookupMsg.KV_REQUEST_ID: "req-1", LookupMsg.KEYS: [b"hX", b"hY", b"hZ"], } ) session.poll() # Dispatch alone must not answer — the parent handle is only valid # during serve_external_requests. assert [ m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupRespMsg.TYPE ] == [] _serve(session, FakeParent()) resps = [ m for m in conn._sent[sent_before:] if m[TYPE_KEY] == LookupRespMsg.TYPE ] assert len(resps) == 1 resp = resps[0] assert resp[LookupRespMsg.KV_REQUEST_ID] == "req-1" assert resp[LookupRespMsg.KEYS] == [b"hX", b"hY", b"hZ"] assert resp[LookupRespMsg.HITS] == [False, False, False] # --------------------------------------------------------------------------- # Server-side handling of inbound LookupMsg (ParentManager-driven) # # poll() only enqueues the LookupMsg; serve_external_requests(parent) # resolves it. Tests follow the poll() → _serve() pattern. # --------------------------------------------------------------------------- def _send_lookup(conn: FakeConnection, kv_request_id: str, keys: list[bytes]): conn.enqueue( { TYPE_KEY: LookupMsg.TYPE, LookupMsg.ROUND_SEQ: 0, LookupMsg.KV_REQUEST_ID: kv_request_id, LookupMsg.KEYS: list(keys), } ) def _lookup_resps(conn: FakeConnection, since: int = 0) -> list[dict]: return [m for m in conn._sent[since:] if m[TYPE_KEY] == LookupRespMsg.TYPE] class TestServerLookupHandling: def test_immediate_hits_create_one_store_job(self): """All-HIT batch: one create_store_job call with all keys, one LookupRespMsg with hits=[True]*N, on_request_finished fires at the end of serve, and `available` is populated for the eventual fetch.""" cb = FakeParent(stored={b"hA": 1, b"hB": 2, b"hC": 3}) session, conn, _ = _make_session() _activate(session, conn) sent_before = len(conn._sent) _send_lookup(conn, "req-1", [b"hA", b"hB", b"hC"]) session.poll() _serve(session, cb) resps = _lookup_resps(conn, sent_before) assert len(resps) == 1 assert resps[0][LookupRespMsg.KEYS] == [b"hA", b"hB", b"hC"] assert resps[0][LookupRespMsg.HITS] == [True, True, True] kinds = [c[0] for c in cb.calls] assert kinds.count("create_store_job") == 1 cs = next(c for c in cb.calls if c[0] == "create_store_job") assert cs[1] == (b"hA", b"hB", b"hC") assert cb.calls[-1][0] == "on_request_finished" # Hits are pinned in outbound state for the upcoming FetchMsg match. assert set(_srv_outbound(session, "req-1").available) == { b"hA", b"hB", b"hC", } def test_all_misses_no_store_job_finish_fires(self): """All-MISS batch: no create_store_job call; one LookupRespMsg with hits=[False]*N; on_request_finished fires at end of serve.""" cb = FakeParent() session, conn, _ = _make_session() _activate(session, conn) sent_before = len(conn._sent) _send_lookup(conn, "req-1", [b"hA", b"hB"]) session.poll() _serve(session, cb) resps = _lookup_resps(conn, sent_before) assert len(resps) == 1 assert resps[0][LookupRespMsg.HITS] == [False, False] assert all(c[0] != "create_store_job" for c in cb.calls) assert cb.calls[-1][0] == "on_request_finished" def test_mixed_hit_miss_pending_defers_response_until_aggregate(self): """HIT/MISS resolutions do not go out on first sight when any key is still HIT_PENDING / RETRY. The lookup parks until every key has settled (or the deadline fires), then one LookupRespMsg carries all keys in wire order.""" cb = FakeParent( stored={b"hA": 1}, pending={b"hB"}, retry={b"hD"}, ) session, conn, _ = _make_session() _activate(session, conn) sent_before = len(conn._sent) _send_lookup(conn, "req-1", [b"hA", b"hB", b"hC", b"hD"]) session.poll() _serve(session, cb) # No LookupRespMsg yet — hB and hD are still pending. assert _lookup_resps(conn, sent_before) == [] # HIT is still pinned immediately so the eventual FetchMsg matches. cs_calls = [c for c in cb.calls if c[0] == "create_store_job"] assert len(cs_calls) == 1 assert cs_calls[0][1] == (b"hA",) # Lookup is parked; on_request_finished not yet called. assert all(c[0] != "on_request_finished" for c in cb.calls) assert len(_srv_lookups(session)) == 1 def test_pending_resolves_then_aggregate_response_fires(self): """A HIT_PENDING key that becomes HIT on a later poll releases the deferred aggregate response: one LookupRespMsg carrying both keys in wire order, and one create_store_job call per HIT (the second HIT is pinned when it resolves, not when the response goes out).""" cb = FakeParent(stored={b"hA": 1}, pending={b"hB"}) session, conn, _ = _make_session() _activate(session, conn) sent_before = len(conn._sent) _send_lookup(conn, "req-1", [b"hA", b"hB"]) session.poll() _serve(session, cb) # No response yet — hB still pending. assert _lookup_resps(conn, sent_before) == [] # Promote hB. cb.pending.discard(b"hB") cb.stored[b"hB"] = 2 # Drive resolver via a second serve_external_requests. _serve(session, cb) resps = _lookup_resps(conn, sent_before) assert len(resps) == 1 assert resps[0][LookupRespMsg.KEYS] == [b"hA", b"hB"] assert resps[0][LookupRespMsg.HITS] == [True, True] cs_calls = [c for c in cb.calls if c[0] == "create_store_job"] assert len(cs_calls) == 2 assert cs_calls[0][1] == (b"hA",) assert cs_calls[1][1] == (b"hB",) # on_request_finished fires once after the aggregate resolve. assert sum(1 for c in cb.calls if c[0] == "on_request_finished") == 1 assert b"hA" in _srv_outbound(session, "req-1").available assert b"hB" in _srv_outbound(session, "req-1").available def test_pending_timeout_replies_miss_no_store_job(self): """A HIT_PENDING key that stays pending past the batch ``deadline`` is force-MISS and never pinned; the deferred aggregate response fires with hits=[False].""" cb = FakeParent(pending={b"hA"}) session, conn, _ = _make_session() _activate(session, conn) sent_before = len(conn._sent) _send_lookup(conn, "req-1", [b"hA"]) session.poll() _serve(session, cb) # Initial serve: nothing immediate, lookup parked, no LookupRespMsg. assert _lookup_resps(conn, sent_before) == [] # Forge the deadline into the past to trigger the timeout branch. lookup = _srv_lookups(session)[0] lookup.deadline = time.monotonic() - 0.1 _serve(session, cb) resps = _lookup_resps(conn, sent_before) assert len(resps) == 1 assert resps[0][LookupRespMsg.KEYS] == [b"hA"] assert resps[0][LookupRespMsg.HITS] == [False] assert all(c[0] != "create_store_job" for c in cb.calls) assert sum(1 for c in cb.calls if c[0] == "on_request_finished") == 1 def test_finish_request_called_per_lookup_msg_not_per_kv_request_id(self): """Two LookupMsgs for the same kv_request_id get distinct ctxs and two on_request_finished calls (one per batch).""" cb = FakeParent(stored={b"hA": 1, b"hB": 2}) session, conn, _ = _make_session() _activate(session, conn) _send_lookup(conn, "req-1", [b"hA"]) session.poll() _serve(session, cb) _send_lookup(conn, "req-1", [b"hB"]) session.poll() _serve(session, cb) finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] assert len(finish_calls) == 2 # Distinct synthetic req_ids assert finish_calls[0][1] != finish_calls[1][1] # Both namespaced under the same kv_request_id assert ":req-1:" in finish_calls[0][1] assert ":req-1:" in finish_calls[1][1] def test_close_returns_open_batch_ctxs_as_failed_serves(self): """Tearing the session down with a parked batch returns the synthetic ctx as a failed serve (no parent handle at teardown) so the manager can release the TieringManager's state on its next serve.""" cb = FakeParent(pending={b"hA"}) session, conn, _ = _make_session() _activate(session, conn) _send_lookup(conn, "req-1", [b"hA"]) session.poll() _serve(session, cb) assert len(_srv_lookups(session)) == 1 assert all(c[0] != "on_request_finished" for c in cb.calls) result = session.close() assert len(result.failed_serves) == 1 assert ":req-1:" in result.failed_serves[0].req_id # close() itself must not call the parent. assert all(c[0] != "on_request_finished" for c in cb.calls) def test_wire_finish_drops_pending_batches_for_kv_request_id(self): """``ServerRole.finish(kv_request_id)`` drops every parked batch whose kv_request_id matches and queues its ctx for the next serve's on_request_finished.""" cb = FakeParent(pending={b"hA", b"hB"}) session, conn, _ = _make_session() _activate(session, conn) _send_lookup(conn, "req-1", [b"hA"]) session.poll() _serve(session, cb) _send_lookup(conn, "req-2", [b"hB"]) session.poll() _serve(session, cb) assert len(_srv_lookups(session)) == 2 session._server.finish("req-1") # req-1 batch dropped from parked lookups; its ctx queued for release. remaining_kv_request_ids = {b.kv_request_id for b in _srv_lookups(session)} assert remaining_kv_request_ids == {"req-2"} queued = session._server._finished_lookup_ctxs assert len(queued) == 1 assert ":req-1:" in queued[0].req_id # The next serve fires on_request_finished exactly once for req-1. _serve(session, cb) finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] assert len(finish_calls) == 1 assert ":req-1:" in finish_calls[0][1] def test_incoming_fetch_drops_pending_lookups_for_kv_request_id(self): """A peer FetchMsg terminates the lookup phase for its id: parked lookups with matching kv_request_id are dropped and their ctx queued for on_request_finished; other kv_request_ids untouched.""" cb = FakeParent(pending={b"hA", b"hB"}) session, conn, _ = _make_session() _activate(session, conn) _send_lookup(conn, "req-1", [b"hA"]) session.poll() _serve(session, cb) _send_lookup(conn, "req-2", [b"hB"]) session.poll() _serve(session, cb) assert len(_srv_lookups(session)) == 2 # Empty FetchMsg: peer signals "lookup phase done" without asking # for any blocks (the all-miss case). conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [], FetchMsg.BLOCK_INDEXES: [], } ) session.poll() remaining_kv_request_ids = {lu.kv_request_id for lu in _srv_lookups(session)} assert remaining_kv_request_ids == {"req-2"} # Dispatch queues the ctx but does not call the parent yet. queued = session._server._finished_lookup_ctxs assert len(queued) == 1 assert ":req-1:" in queued[0].req_id # The next serve fires on_request_finished exactly once for req-1. _serve(session, cb) finish_calls = [c for c in cb.calls if c[0] == "on_request_finished"] assert len(finish_calls) == 1 assert ":req-1:" in finish_calls[0][1] def test_lookup_then_fetch_round_trip_emits_store_result(self): """End-to-end: lookup pins primary slots → fetch matches them → NIXL transfer completes → StoreResult surfaces with the create_store_job's job_id (the engine releases the pin).""" cb = FakeParent(stored={b"hA": 7, b"hB": 8}) session, conn, transport = _make_session() _activate(session, conn) _send_lookup(conn, "req-1", [b"hA", b"hB"]) session.poll() _serve(session, cb) cs = next(c for c in cb.calls if c[0] == "create_store_job") # FakeParent issues monotonic job_ids starting at 1000. expected_job_id = 1000 # Consumer issues FetchMsg on the resolved hits. conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"hA", b"hB"], FetchMsg.BLOCK_INDEXES: [20, 21], } ) session.poll() # NIXL write_blocks called with our pinned local block_ids. assert len(transport._transfers) == 1 _, (_peer, local, remote) = next(iter(transport._transfers.items())) assert local == [7, 8] assert remote == [20, 21] # Drive the transport completion. transport._poll_done.append(0) result = session.poll() store_results = [s for s in result.stores if s.success] assert any(s.job_id == expected_job_id for s in store_results) # Sanity: kv mention in synthetic ctx. assert cs[2].startswith("p2p:") # --------------------------------------------------------------------------- # Server-role flows # --------------------------------------------------------------------------- class TestServerFlows: def test_store_then_fetch_matches(self): """Blocks stored before fetch demand are matched on demand arrival.""" session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1", b"k2"], [0, 1], job_id=1) conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) session.poll() assert len(transport._transfers) == 1 _, (peer, local, remote) = next(iter(transport._transfers.items())) assert local == [0, 1] assert remote == [10, 11] def test_fetch_then_store_matches(self): """Fetch demand registered before store; store fulfills it.""" session, conn, transport = _make_session() _activate(session, conn) conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) session.poll() assert len(transport._transfers) == 0 session.add_stored_blocks("req-1", [b"k1"], [3], job_id=1) assert len(transport._transfers) == 1 def test_transfer_completion_emits_store_result_and_done(self): """Completed transfer reports StoreResult and sends TransferDoneMsg.""" session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) session.poll() tid = next(iter(transport._transfers)) transport._poll_done.append(tid) stores = session.poll().stores assert StoreResult(job_id=1, success=True) in stores assert any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent) def test_abort_fetch_replies_with_ack(self): session, conn, _ = _make_session() _activate(session, conn) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" assert _srv_abort_started(session, "req-1") is None def test_abort_fetch_defers_ack_when_cancel_pending(self): """If cancel(mode='wait') reports still-inflight tids, the ack is deferred and the abort is parked (abort_started_at set).""" session, conn, transport = _make_session() _activate(session, conn) # Seed an inflight transfer for req-1 that the transport pretends # cannot be canceled yet. tid = 42 session._server._inflight_add( tid, _InflightXfer( kv_request_id="req-1", block_count=1, job_ids={1}, round=_OutboundRequestState(inflight=1), ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) assert _srv_abort_started(session, "req-1") is not None # First attempt happens inside _on_abort_fetch; the per-tick # drain runs again at the end of poll() — both are wait-mode. assert all(mode == "wait" for _, mode in transport._cancel_calls) assert tid in session._server._inflight # still tracked def test_abort_fetch_acks_after_drain(self): """Once the transport reports the tid as DONE the parked abort completes and AbortAckMsg is sent.""" session, conn, transport = _make_session() _activate(session, conn) tid = 42 session._server._inflight_add( tid, _InflightXfer( kv_request_id="req-1", block_count=1, job_ids={1}, round=_OutboundRequestState(inflight=1), ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() assert _srv_abort_started(session, "req-1") is not None # Backend finishes draining: transport.poll() will return tid as # DONE, and the next cancel(mode='wait') call sees it's gone. transport._cancel_still_inflight.discard(tid) transport._poll_done.append(tid) session.poll() ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" assert _srv_abort_started(session, "req-1") is None assert tid not in session._server._inflight def test_abort_fetch_force_cancels_after_timeout(self): """If wait-mode never drains, the deadline forces immediate cancel and an ack is still sent.""" session, conn, transport = _make_session() _activate(session, conn) tid = 42 session._server._inflight_add( tid, _InflightXfer( kv_request_id="req-1", block_count=1, job_ids={1}, round=_OutboundRequestState(inflight=1), ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() assert _srv_abort_started(session, "req-1") is not None # Backdate past the drain deadline. session._server._pending_aborts[("req-1", 0)] = ( time.monotonic() - _CANCEL_DRAIN_TIMEOUT_S - 1.0 ) # Even if the transport still claims it can't cancel, the # session must force-pop and ack. transport._cancel_calls.clear() session.poll() ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE) assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1" assert _srv_abort_started(session, "req-1") is None assert tid not in session._server._inflight assert ([tid], "immediate") in transport._cancel_calls def test_abort_fetch_idempotent_while_draining(self): """Receiving AbortFetchMsg twice for the same kv_request_id keeps a single pending entry and produces a single ack.""" session, conn, transport = _make_session() _activate(session, conn) tid = 42 session._server._inflight_add( tid, _InflightXfer( kv_request_id="req-1", block_count=1, job_ids={1}, round=_OutboundRequestState(inflight=1), ), ) transport._cancel_still_inflight.add(tid) conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() first_started_at = _srv_abort_started(session, "req-1") # Second AbortFetchMsg for the same kv_request_id while still # draining must not reset the deadline. conn.enqueue( { TYPE_KEY: AbortFetchMsg.TYPE, AbortFetchMsg.ROUND_SEQ: 0, AbortFetchMsg.KV_REQUEST_ID: "req-1", } ) session.poll() assert _srv_abort_started(session, "req-1") == first_started_at # Now let the drain succeed and confirm exactly one ack ever. transport._cancel_still_inflight.discard(tid) transport._poll_done.append(tid) session.poll() acks = [m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE] assert len(acks) == 1 assert acks[0][AbortAckMsg.KV_REQUEST_ID] == "req-1" def test_store_timeout(self): session, conn, _ = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) # Backdate. session._server._store_jobs[1] = time.monotonic() - 60.0 stores = session.poll().stores assert StoreResult(job_id=1, success=False) in stores def test_store_timeout_then_late_completion_no_duplicate(self): """A job timed out by _timeout_pending_store_jobs must not also emit a contradictory StoreResult(success=True) when the transport later reports the same transfer as done.""" session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) session.poll() tid = next(iter(transport._transfers)) # Backdate the store job so the next poll times it out. session._server._store_jobs[1] = time.monotonic() - 60.0 stores = session.poll().stores assert StoreResult(job_id=1, success=False) in stores assert StoreResult(job_id=1, success=True) not in stores # Transport later reports the same transfer as done — must not # emit a second (contradictory) StoreResult for job_id=1. transport._poll_done.append(tid) stores = session.poll().stores assert all(s.job_id != 1 for s in stores), ( f"unexpected duplicate StoreResult after timeout: {stores}" ) def test_store_timeout_then_late_failure_no_duplicate(self): """Symmetric guard: a timed-out job must not also emit a second StoreResult(success=False) when the transport later reports the same transfer as failed.""" session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) session.poll() tid = next(iter(transport._transfers)) session._server._store_jobs[1] = time.monotonic() - 60.0 stores = session.poll().stores assert [s for s in stores if s.job_id == 1] == [ StoreResult(job_id=1, success=False) ] transport._poll_failed.append(tid) stores = session.poll().stores assert all(s.job_id != 1 for s in stores), ( f"unexpected duplicate StoreResult after timeout: {stores}" ) # --------------------------------------------------------------------------- # finish_request server-role early-fail flow # --------------------------------------------------------------------------- class TestFinishRequestServerSide: def _last_transfer_done(self, conn: FakeConnection) -> dict | None: for msg in reversed(conn._sent): if msg[TYPE_KEY] == TransferDoneMsg.TYPE: return msg return None def test_no_inflight_unmatched_demand_sends_failure(self): """finish_request with unmatched demand and no inflight -> immediate TransferDoneMsg(success=False); _outbound cleared.""" session, conn, _ = _make_session() _activate(session, conn) # Decoder demanded a block we never stored. conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [5], } ) session.poll() assert _srv_outbound(session, "req-1") is not None session.finish_request("req-1") msg = self._last_transfer_done(conn) assert msg is not None assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert msg[TransferDoneMsg.SUCCESS] is False assert _srv_outbound(session, "req-1") is None def test_with_inflight_defers_then_fires_on_last_transfer(self): """finish_request with inflight defers; last transfer fires the early-fail message and clears _outbound.""" session, conn, transport = _make_session() _activate(session, conn) # Demand 2 blocks; we store 1 (kicks one inflight transfer). conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) session.poll() session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) assert len(transport._transfers) == 1 # finish_request while inflight: no early-fail yet. before = len(conn._sent) session.finish_request("req-1") assert len(conn._sent) == before assert _srv_outbound(session, "req-1") is not None assert _srv_outbound(session, "req-1").finishing # Last inflight settles -> early-fail fires. tid = next(iter(transport._transfers)) transport._poll_done.append(tid) session.poll() msg = self._last_transfer_done(conn) assert msg is not None assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert msg[TransferDoneMsg.SUCCESS] is False assert _srv_outbound(session, "req-1") is None def test_full_demand_satisfied_still_sends_success(self): """finish_request must not override a fully-satisfied transfer: when remaining hits 0, success=True still fires.""" session, conn, transport = _make_session() _activate(session, conn) conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) session.poll() session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) # Mark finishing (e.g., on_request_finished racing with the last # store) — but all demand is satisfied. session.finish_request("req-1") tid = next(iter(transport._transfers)) transport._poll_done.append(tid) session.poll() msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE) assert msg[TransferDoneMsg.SUCCESS] is True def test_prefiller_first_finish_before_fetch(self): """Prefiller-first: finish_request runs before the decoder's fetch arrives. State is held until fetch, then finalized — success=True if all demand was matched against available blocks, else success=False.""" # Case A: all demand satisfied by available blocks. session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) # finish_request first — no demand received yet -> defer. session.finish_request("req-1") assert _srv_outbound(session, "req-1") is not None # Fetch arrives now: demand fully satisfied by available. conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) session.poll() # The transfer was inflight; finalize it. tid = next(iter(transport._transfers)) transport._poll_done.append(tid) session.poll() msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE) assert msg[TransferDoneMsg.SUCCESS] is True # Case B: demand exceeds available -> early-fail fires from fetch. session, conn, transport = _make_session() _activate(session, conn) session.add_stored_blocks("req-2", [b"k1"], [0], job_id=2) session.finish_request("req-2") conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-2", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) session.poll() # Inflight for k1 still in flight; nothing yet for the early-fail # — the same code path will fire from _collect_store_results. tid = next(iter(transport._transfers)) transport._poll_done.append(tid) session.poll() msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE) assert msg[TransferDoneMsg.SUCCESS] is False assert _srv_outbound(session, "req-2") is None def test_unknown_request_is_noop(self): session, conn, _ = _make_session() _activate(session, conn) before = len(conn._sent) session.finish_request("never-existed") assert len(conn._sent) == before def test_finish_request_no_inflight_emits_store_failure(self): """finish_request with a stored-but-unmatched job and no inflight -> TransferDoneMsg(success=False) AND deferred StoreResult(success=False) for the submit_store'd job, instead of the 30s _STORE_TIMEOUT_S path.""" session, conn, _ = _make_session() _activate(session, conn) # Decoder demanded b"demand"; we never stored it. conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"demand"], FetchMsg.BLOCK_INDEXES: [5], } ) session.poll() # We did submit_store a different block — goes to available, never # matches demand. Without the shortcut, job 42 sits in _store_jobs # for _STORE_TIMEOUT_S. session.add_stored_blocks("req-1", [b"unrelated"], [0], job_id=42) assert _srv_outbound(session, "req-1").pending_job_ids == {42} session.finish_request("req-1") # Peer notified immediately with success=False (remaining > 0). msg = self._last_transfer_done(conn) assert msg is not None assert msg[TransferDoneMsg.SUCCESS] is False assert _srv_outbound(session, "req-1") is None # Local store job surfaces on the next poll, success=False. stores = session.poll().stores assert StoreResult(job_id=42, success=False) in stores assert 42 not in session._server._store_jobs def test_finish_request_remaining_zero_emits_success_via_inflight(self): """Deferred-via-inflight path: finish_request with inflight, last transfer drains remaining to 0 -> TransferDoneMsg(success=True) AND StoreResult(success=True).""" session, conn, transport = _make_session() _activate(session, conn) conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) session.poll() session.add_stored_blocks("req-1", [b"k1"], [0], job_id=7) # finish_request races with the inflight transfer. session.finish_request("req-1") assert _srv_outbound(session, "req-1") is not None # deferred # Last inflight completes -> _finalize_outbound(success=True) fires. tid = next(iter(transport._transfers)) transport._poll_done.append(tid) stores = session.poll().stores msg = self._last_transfer_done(conn) assert msg is not None assert msg[TransferDoneMsg.SUCCESS] is True assert StoreResult(job_id=7, success=True) in stores assert _srv_outbound(session, "req-1") is None def test_write_blocks_failure_finalizes_with_failure(self): """write_blocks returning None must not leave the request hanging. The matched blocks are gone from req.demanded but no inflight will satisfy them, so remaining > 0 forever. Setting finishing and calling _finalize_outbound(success=False) immediately (no other inflight) tells the peer + emits StoreResult(success=False) without waiting on _STORE_TIMEOUT_S or _LOAD_TIMEOUT_S. """ session, conn, transport = _make_session() _activate(session, conn) # Decoder demands b"k1". conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [10], } ) session.poll() # Force write_blocks to fail on the next call. transport.write_blocks = lambda *a, **kw: None # type: ignore[assignment] session.add_stored_blocks("req-1", [b"k1"], [0], job_id=42) # Outbound was finalized immediately (no other inflight). assert _srv_outbound(session, "req-1") is None # Peer notified with success=False. msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE) assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert msg[TransferDoneMsg.SUCCESS] is False # Local store job surfaces on the next poll. stores = session.poll().stores assert StoreResult(job_id=42, success=False) in stores assert 42 not in session._server._store_jobs def test_partial_match_completes_in_two_rounds(self): """Peer demand for [k1, k2, k3]; first round only k1 is available, second round adds k2 and k3. Each round transfers what's matched and the request finalizes with success once remaining hits zero. """ session, conn, transport = _make_session() _activate(session, conn) # Peer fetches three blocks. conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2", b"k3"], FetchMsg.BLOCK_INDEXES: [10, 11, 12], } ) session.poll() # Demand registered, no matches yet. assert session._server._inflight == {} outbound = _srv_outbound(session, "req-1") assert outbound.remaining == 3 assert set(outbound.demanded.keys()) == {b"k1", b"k2", b"k3"} # Round 1: only k1 is stored locally. session.add_stored_blocks("req-1", [b"k1"], [0], job_id=100) # One inflight transfer for k1. assert len(session._server._inflight) == 1 tid_1 = next(iter(session._server._inflight)) assert transport._transfers[tid_1][1] == [0] assert transport._transfers[tid_1][2] == [10] # k2 and k3 still demanded. assert set(outbound.demanded.keys()) == {b"k2", b"k3"} # Transfer 1 completes. transport._poll_done.append(tid_1) stores = session.poll().stores assert StoreResult(job_id=100, success=True) in stores assert session._server._inflight == {} assert outbound.remaining == 2 # Not yet finalized — still 2 blocks demanded. assert _srv_outbound(session, "req-1") is not None # Round 2: k2 and k3 arrive together. session.add_stored_blocks("req-1", [b"k2", b"k3"], [1, 2], job_id=200) assert len(session._server._inflight) == 1 tid_2 = next(iter(session._server._inflight)) assert tid_2 != tid_1 assert sorted(transport._transfers[tid_2][1]) == [1, 2] # Transfer 2 completes — request now fully satisfied. transport._poll_done.append(tid_2) stores = session.poll().stores assert StoreResult(job_id=200, success=True) in stores # _finalize_outbound fired — request gone, peer notified with success. assert _srv_outbound(session, "req-1") is None done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE) assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert done[TransferDoneMsg.SUCCESS] is True def test_write_blocks_failure_finalizes_after_last_inflight_completes(self): """write_blocks returns None on a SECOND match while a first transfer is still inflight. The request should NOT finalize until the inflight completes, then the elif branch in collect_results (``finishing and not _has_inflight_for(...)``) finalizes it as failure. """ session, conn, transport = _make_session() _activate(session, conn) # Peer demands two blocks. conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [10, 11], } ) session.poll() # Round 1: k1 transfers cleanly. session.add_stored_blocks("req-1", [b"k1"], [0], job_id=100) assert len(session._server._inflight) == 1 tid_1 = next(iter(session._server._inflight)) outbound = _srv_outbound(session, "req-1") assert outbound.remaining == 2 # decrement happens on completion assert outbound.finishing is False # Round 2: write_blocks fails for k2 while transfer_1 is still inflight. transport.write_blocks = lambda *a, **kw: None # type: ignore[assignment] session.add_stored_blocks("req-1", [b"k2"], [1], job_id=200) # No new transfer was registered. assert list(session._server._inflight.keys()) == [tid_1] # Marked finishing, but NOT finalized yet (transfer_1 still inflight). assert outbound.finishing is True assert _srv_outbound(session, "req-1") is not None done_msgs = [m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE] assert done_msgs == [] # Transfer 1 completes — now ``_has_inflight_for("req-1")`` is False # and the elif branch in collect_results fires _finalize(success=False). # The k1 success result is direct; the k2 failure result is queued # in _pending_store_results and surfaces on the NEXT poll. transport._poll_done.append(tid_1) stores_first = session.poll().stores assert StoreResult(job_id=100, success=True) in stores_first # Outbound state cleaned up; peer notified with success=False. assert _srv_outbound(session, "req-1") is None done = next(m for m in conn._sent if m.get(TYPE_KEY) == TransferDoneMsg.TYPE) assert done[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert done[TransferDoneMsg.SUCCESS] is False # Next poll drains the queued failure. stores_second = session.poll().stores assert StoreResult(job_id=200, success=False) in stores_second # --------------------------------------------------------------------------- # Bidirectional — the case the unification is meant to fix # --------------------------------------------------------------------------- class TestBidirectional: def test_session_handles_both_roles_concurrently(self): """Single session simultaneously serves a fetch and completes a load. This is the regression test for the unification: with the old split design the inbound FetchMsg would be dispatched to a client-only session and dropped (or a server-only session would miss the TransferDoneMsg). One unified session handles both. """ session, conn, transport = _make_session() _activate(session, conn) # Server role: the peer fetches a block from us. session.add_stored_blocks("req-srv", [b"served"], [0], job_id=100) conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-srv", FetchMsg.KEYS: [b"served"], FetchMsg.BLOCK_INDEXES: [7], } ) # Client role: we ask the peer for a different block. session.request_blocks( job_id=200, kv_request_id="req-cli", keys=[b"loaded"], block_ids=[3] ) # Both flows progress in the same poll. session.poll() # Server side: write_blocks was submitted. assert len(transport._transfers) == 1 # Client side: the lookup was sent. assert any( m[TYPE_KEY] == FetchMsg.TYPE and m[FetchMsg.KV_REQUEST_ID] == "req-cli" for m in conn._sent ) # Peer now signals: server-side transfer completes AND a # TransferDoneMsg arrives for our client-side request, all in # one batch on the same connection. tid = next(iter(transport._transfers)) transport._poll_done.append(tid) conn.enqueue( { TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-cli", TransferDoneMsg.SUCCESS: True, } ) result_ = session.poll() loads = result_.loads stores = result_.stores assert LoadResult(job_id=200, kv_request_id="req-cli", success=True) in loads assert StoreResult(job_id=100, success=True) in stores # --------------------------------------------------------------------------- # Pending sessions (no connection yet) # --------------------------------------------------------------------------- class TestPendingSession: def test_pending_session_buffers_stored_blocks(self): """Pending session accepts add_stored_blocks but cannot send.""" transport = FakeDataTransport() session = P2PSession( peer_id="peer:8000", local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) result_ = session.poll() loads = result_.loads stores = result_.stores assert loads == [] assert stores == [] assert not session.connected assert session.alive def test_attach_connection_sends_connect(self): """attach_connection triggers our ConnectMsg send.""" transport = FakeDataTransport() session = P2PSession( peer_id="peer:8000", local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) conn = FakeConnection(peer_id="peer:8000") session.attach_connection(conn) # type: ignore[arg-type] assert conn._sent assert conn._sent[0][TYPE_KEY] == ConnectMsg.TYPE def test_attach_connection_twice_raises(self): """attach_connection on an already-connected session raises.""" session, conn, _ = _make_session() with pytest.raises(ValueError, match="already connected"): session.attach_connection(FakeConnection()) # type: ignore[arg-type] def test_pending_close_returns_pending_stores(self): """Closing a pending session reports buffered stores as failed.""" transport = FakeDataTransport() session = P2PSession( peer_id="peer:8000", local_id="local:9000", transport=transport, # type: ignore[arg-type] local_block_len=4096, local_hash_seed=_DEFAULT_HASH_SEED, conn=None, ) session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1) session.add_stored_blocks("req-2", [b"k2"], [1], job_id=2) result = session.close() assert result.failed_jobs == [] assert result.failed_req_ids == [] assert set(result.failed_stores) == {1, 2} assert result.failed_serves == [] # --------------------------------------------------------------------------- # Disconnect # --------------------------------------------------------------------------- class TestDisconnect: def test_disconnect_marks_session_dead(self): session, conn, _ = _make_session() _activate(session, conn) conn.enqueue({TYPE_KEY: DisconnectMsg.TYPE}) session.poll() assert not session.alive def test_close_returns_pending_loads_and_stores(self): session, conn, _ = _make_session() _activate(session, conn) session.request_blocks(1, "req-1", [b"k"], [0]) session.request_blocks(2, "req-2", [b"k"], [0]) session.add_stored_blocks("req-srv", [b"k"], [0], job_id=10) result = session.close() assert set(result.failed_jobs) == {1, 2} assert set(result.failed_req_ids) == {"req-1", "req-2"} assert set(result.failed_stores) == {10} assert result.failed_serves == [] def test_send_failure_marks_connection_dead(self): """A raising send must mark the connection dead, not silently drop the message — otherwise the session lingers alive, is never reaped, and in-flight lookups/loads toward the dead peer hang forever.""" session, conn, _ = _make_session() _activate(session, conn) assert session.alive conn.fail_send = True # request_blocks flushes a FetchMsg synchronously via _do_send. session.request_blocks(1, "req-1", [b"k"], [0]) assert not session.alive def test_close_surfaces_inflight_lookups(self): """close() reports kv_request_ids whose symmetric-P2P probe is still unresolved; resolved probes are not reported (their answer is in).""" session, conn, _ = _make_session() _activate(session, conn) session.register_lookup("req-hit", b"hA") session.register_lookup("req-inflight", b"hB") session.flush_pending_lookups() # Only req-hit is answered; req-inflight stays in flight. conn.enqueue( { TYPE_KEY: LookupRespMsg.TYPE, LookupRespMsg.KV_REQUEST_ID: "req-hit", LookupRespMsg.KEYS: [b"hA"], LookupRespMsg.HITS: [True], } ) session.poll() result = session.close() assert result.failed_jobs == [] assert result.failed_req_ids == ["req-inflight"] # --------------------------------------------------------------------------- # Adversarial / malformed messages # --------------------------------------------------------------------------- class TestAdversarial: def test_unknown_message_type_logged(self): session, conn, _ = _make_session() _activate(session, conn) conn.enqueue({TYPE_KEY: "evil_command"}) result_ = session.poll() loads = result_.loads stores = result_.stores assert loads == [] assert stores == [] def test_empty_message(self): session, conn, _ = _make_session() _activate(session, conn) conn.enqueue({}) result_ = session.poll() loads = result_.loads stores = result_.stores assert loads == [] assert stores == [] def test_non_dict_message(self): session, conn, _ = _make_session() _activate(session, conn) conn._inbox.append(42) # type: ignore[arg-type] result_ = session.poll() loads = result_.loads stores = result_.stores assert loads == [] assert stores == [] def test_fetch_mismatched_lengths(self): session, conn, transport = _make_session() _activate(session, conn) conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-bad", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], } ) session.poll() assert len(transport._transfers) == 0 # Protocol violation: session disconnects immediately so the peer # can't keep wedging us with malformed traffic. assert not session.alive assert any(m[TYPE_KEY] == DisconnectMsg.TYPE for m in conn._sent) def test_transfer_done_missing_kv_request_id(self): session, conn, _ = _make_session() _activate(session, conn) conn.enqueue({TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.SUCCESS: True}) loads = session.poll().loads assert loads == [] assert not session.alive assert any(m[TYPE_KEY] == DisconnectMsg.TYPE for m in conn._sent) def test_duplicate_connect_ack(self): session, conn, _ = _make_session() _activate(session, conn) conn.enqueue({TYPE_KEY: ConnectAckMsg.TYPE, ConnectAckMsg.PEER_ID: "peer:8000"}) session.poll() assert session.ready class TestDispatchErrorHandling: """Errors raised by message handlers split into two classes: - Protocol-contract violations from the peer (ValueError) → disconnect on the first occurrence; retrying won't help and may corrupt state. - Anything else is treated as an internal bug: log loudly, count, and only disconnect once errors arrive in a tight burst. A successful dispatch in between resets the counter. """ def test_value_error_disconnects_on_first_occurrence(self): """A FetchMsg that fails validate() raises ValueError and must terminate the session immediately, with a DisconnectMsg sent.""" session, conn, _ = _make_session() _activate(session, conn) # length mismatch → FetchMsg.validate raises ValueError conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-bad", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [1], } ) session.poll() assert not session.alive assert any(m[TYPE_KEY] == DisconnectMsg.TYPE for m in conn._sent) def test_transfer_done_missing_field_disconnects(self): """Same contract for a malformed TransferDoneMsg from the peer.""" session, conn, _ = _make_session() _activate(session, conn) conn.enqueue({TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.SUCCESS: True}) session.poll() assert not session.alive def test_internal_error_does_not_disconnect_once(self): """A non-ValueError raised by a handler is treated as an internal bug: counter increments, session stays alive on a single hit.""" session, conn, _ = _make_session() _activate(session, conn) def _boom(*args, **kwargs): raise RuntimeError("simulated internal bug") session._server.on_fetch = _boom # type: ignore[assignment] conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) session.poll() assert session.alive assert session._dispatch_error_count == 1 def test_internal_error_threshold_disconnects(self): """Once consecutive non-protocol errors hit the threshold, the session tears down via _protocol_error.""" session, conn, _ = _make_session() _activate(session, conn) def _boom(*args, **kwargs): raise RuntimeError("simulated internal bug") session._server.on_fetch = _boom # type: ignore[assignment] for _ in range(_MAX_CONSECUTIVE_DISPATCH_ERRORS): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) session.poll() assert not session.alive assert any(m[TYPE_KEY] == DisconnectMsg.TYPE for m in conn._sent) def test_internal_error_counter_resets_on_success(self): """A successful dispatch between errors prevents the threshold from being reached.""" session, conn, _ = _make_session() _activate(session, conn) original_on_fetch = session._server.on_fetch def _boom(*args, **kwargs): raise RuntimeError("simulated internal bug") # Alternate (boom, success) (_MAX-1) times: counter rises to 1 # then resets to 0 each cycle, never reaching the threshold. for _ in range(_MAX_CONSECUTIVE_DISPATCH_ERRORS - 1): session._server.on_fetch = _boom # type: ignore[assignment] conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1"], FetchMsg.BLOCK_INDEXES: [0], } ) session.poll() session._server.on_fetch = original_on_fetch # type: ignore[assignment] # A benign no-op message (unknown type) dispatches cleanly # and resets the consecutive-error counter. conn.enqueue({TYPE_KEY: "unknown_for_test"}) session.poll() assert session.alive assert session._dispatch_error_count == 0 class TestInflightPerReqInvariant: """Per-request `inflight_tids` is the O(1) replacement for the previous O(N) scan in `_has_inflight_for`. These tests check that every mutation site keeps the set in sync with `_inflight` and that the lookup is correct under high fan-out. """ def test_invariant_holds_through_lifecycle(self): """Run a full submit→complete sequence for two concurrent kv_request_ids and assert the counter matches `_inflight` at every observable step, including the empty-after-finish case.""" session, conn, transport = _make_session() _activate(session, conn) def _invariant_holds() -> bool: return _srv_total_inflight(session) == len(session._server._inflight) assert _invariant_holds() # Two requests, two blocks each, all dispatched in one batch. session.add_stored_blocks("req-A", [b"a1", b"a2"], [0, 1], job_id=10) session.add_stored_blocks("req-B", [b"b1", b"b2"], [2, 3], job_id=11) for kv_id, keys, indexes in ( ("req-A", [b"a1", b"a2"], [100, 101]), ("req-B", [b"b1", b"b2"], [102, 103]), ): conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: kv_id, FetchMsg.KEYS: keys, FetchMsg.BLOCK_INDEXES: indexes, } ) session.poll() assert _invariant_holds() assert session._server._has_inflight_for("req-A") assert session._server._has_inflight_for("req-B") assert not session._server._has_inflight_for("req-C") # Complete req-A's transfer first; req-B should still be inflight. a_tids = [ tid for tid, x in session._server._inflight.items() if x.kv_request_id == "req-A" ] for tid in a_tids: transport._poll_done.append(tid) session.poll() assert _invariant_holds() assert not session._server._has_inflight_for("req-A") assert _srv_inflight_count(session, "req-A") == 0 # entry drained assert session._server._has_inflight_for("req-B") # Complete req-B; counter must drain to empty. b_tids = [ tid for tid, x in session._server._inflight.items() if x.kv_request_id == "req-B" ] for tid in b_tids: transport._poll_done.append(tid) session.poll() assert _invariant_holds() assert session._server._inflight == {} assert _srv_total_inflight(session) == 0 def test_has_inflight_for_correct_with_many_requests(self): """Populate many inflight xfers across many ids; lookup must match the actual presence in `_inflight` for both hits and misses. The whole point of the counter is that this lookup is constant-time, but we assert correctness, not timing.""" session, _, _ = _make_session() for kv_id_idx in range(100): kv_id = f"req-{kv_id_idx}" for j in range(10): tid = kv_id_idx * 10 + j session._server._inflight_add( tid, _InflightXfer( kv_request_id=kv_id, block_count=1, job_ids={tid}, round=_OutboundRequestState(inflight=1), ), ) assert _srv_total_inflight(session) == len(session._server._inflight) assert session._server._has_inflight_for("req-0") assert session._server._has_inflight_for("req-99") assert not session._server._has_inflight_for("req-missing") # Drain all entries for req-50 and confirm the entry disappears. tids_50 = [ tid for tid, x in session._server._inflight.items() if x.kv_request_id == "req-50" ] for tid in tids_50: session._server._inflight_pop(tid) assert _srv_inflight_count(session, "req-50") == 0 assert not session._server._has_inflight_for("req-50") # Other ids unaffected. assert session._server._has_inflight_for("req-49") # --------------------------------------------------------------------------- # Protocol validation tests (unchanged from the prior file) # --------------------------------------------------------------------------- class TestConnectMsgValidation: def _valid_msg(self) -> dict: return { TYPE_KEY: ConnectMsg.TYPE, ConnectMsg.PEER_ID: "peer:1", ConnectMsg.AGENT_METADATA: b"meta", ConnectMsg.BASE_ADDR: 0x1000, ConnectMsg.NUM_BLOCKS: 8, ConnectMsg.BLOCK_LEN: 4096, ConnectMsg.HASH_SEED: "0", } def test_valid_message_passes(self): ConnectMsg.validate(self._valid_msg()) def test_missing_peer_id(self): msg = self._valid_msg() del msg[ConnectMsg.PEER_ID] with pytest.raises(ValueError, match="peer_id"): ConnectMsg.validate(msg) def test_missing_agent_metadata(self): msg = self._valid_msg() del msg[ConnectMsg.AGENT_METADATA] with pytest.raises(ValueError, match="agent_metadata"): ConnectMsg.validate(msg) def test_agent_metadata_wrong_type(self): msg = self._valid_msg() msg[ConnectMsg.AGENT_METADATA] = "not bytes" with pytest.raises(ValueError, match="agent_metadata"): ConnectMsg.validate(msg) def test_base_addr_negative(self): msg = self._valid_msg() msg[ConnectMsg.BASE_ADDR] = -1 with pytest.raises(ValueError, match="base_addr"): ConnectMsg.validate(msg) def test_num_blocks_zero(self): msg = self._valid_msg() msg[ConnectMsg.NUM_BLOCKS] = 0 with pytest.raises(ValueError, match="num_blocks"): ConnectMsg.validate(msg) def test_block_len_zero(self): msg = self._valid_msg() msg[ConnectMsg.BLOCK_LEN] = 0 with pytest.raises(ValueError, match="block_len"): ConnectMsg.validate(msg) def test_missing_hash_seed(self): msg = self._valid_msg() del msg[ConnectMsg.HASH_SEED] with pytest.raises(ValueError, match="hash_seed"): ConnectMsg.validate(msg) def test_hash_seed_wrong_type(self): msg = self._valid_msg() msg[ConnectMsg.HASH_SEED] = 12345 # int, not str with pytest.raises(ValueError, match="hash_seed"): ConnectMsg.validate(msg) class TestFetchMsgValidation: def _valid_msg(self) -> dict: return { TYPE_KEY: FetchMsg.TYPE, FetchMsg.ROUND_SEQ: 0, FetchMsg.KV_REQUEST_ID: "req-1", FetchMsg.KEYS: [b"k1", b"k2"], FetchMsg.BLOCK_INDEXES: [0, 1], } def test_valid_message_passes(self): FetchMsg.validate(self._valid_msg()) def test_length_mismatch(self): msg = self._valid_msg() msg[FetchMsg.BLOCK_INDEXES] = [0] with pytest.raises(ValueError, match="length mismatch"): FetchMsg.validate(msg) def test_negative_index(self): msg = self._valid_msg() msg[FetchMsg.BLOCK_INDEXES] = [0, -1] with pytest.raises(ValueError, match="invalid index"): FetchMsg.validate(msg) class TestTransferDoneMsgValidation: def test_valid_message_passes(self): msg = { TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: True, } TransferDoneMsg.validate(msg) def test_success_wrong_type(self): msg = { TYPE_KEY: TransferDoneMsg.TYPE, TransferDoneMsg.ROUND_SEQ: 0, TransferDoneMsg.KV_REQUEST_ID: "req-1", TransferDoneMsg.SUCCESS: 1, } with pytest.raises(ValueError, match="success"): TransferDoneMsg.validate(msg)