# SPDX-License-Identifier: Apache-2.0
"""
Unit tests for PrefetchController.

Tests verify the end-to-end prefetch flow: submit request → lookup in L2 →
compute prefix-trimmed load plan → reserve L1 buffers → load from L2 →
transition to read-locked → report prefix hits.

Uses a real L1Manager and MockL2Adapter (with debug methods) to exercise
the full integration without mocking internals.
"""

# Standard
import time

# Third Party
import pytest
import torch

# First Party
from lmcache.native_storage_ops import Bitmap
from lmcache.v1.distributed.api import (
    MemoryLayoutDesc,
    ObjectKey,
    PrefetchMode,
    TrimPolicy,
)
from lmcache.v1.distributed.config import L1ManagerConfig, L1MemoryManagerConfig
from lmcache.v1.distributed.error import L1Error
from lmcache.v1.distributed.l1_manager import L1Manager
from lmcache.v1.distributed.l2_adapters.fault_inject_l2_adapter import (
    FaultInjectL2Adapter,
)
from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import (
    MockL2Adapter,
    MockL2AdapterConfig,
)
from lmcache.v1.distributed.storage_controllers.prefetch_controller import (
    PrefetchController,
    build_trim_mask,
    merge_bitmaps,
)
from lmcache.v1.distributed.storage_controllers.prefetch_policy import (
    DefaultPrefetchPolicy,
)
from lmcache.v1.distributed.storage_controllers.store_policy import (
    AdapterDescriptor,
)
from lmcache.v1.memory_management import MemoryObjMetadata, TensorMemoryObj

# Skip all tests in this module if CUDA is not available
pytestmark = pytest.mark.skipif(
    not torch.cuda.is_available(), reason="CUDA is not available"
)


# =============================================================================
# Helpers
# =============================================================================


def make_object_key(chunk_id: int) -> ObjectKey:
    """Create a test ObjectKey with the given chunk ID."""
    return ObjectKey(
        chunk_hash=ObjectKey.IntHash2Bytes(chunk_id),
        model_name="test_model",
        kv_rank=0,
    )


def make_layout() -> MemoryLayoutDesc:
    """Create a small MemoryLayoutDesc for testing."""
    return MemoryLayoutDesc(
        shapes=[torch.Size([100, 2, 512])],
        dtypes=[torch.bfloat16],
    )


def should_use_lazy_alloc() -> bool:
    return torch.cuda.is_available()


def wait_for_condition(
    predicate,
    timeout: float = 5.0,
    poll_interval: float = 0.05,
) -> bool:
    """Poll until a predicate returns True or timeout."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return True
        time.sleep(poll_interval)
    return False


def wait_for_prefetch_result(
    ctrl: PrefetchController,
    req_id: int,
    timeout: float = 5.0,
    poll_interval: float = 0.05,
) -> int | None:
    """Poll query_prefetch_result until it returns a non-None value."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        result = ctrl.query_prefetch_result(req_id)
        if result is not None:
            return result.count_leading_ones()
        time.sleep(poll_interval)
    return None


def wait_for_lookup_result(
    ctrl: PrefetchController,
    req_id: int,
    timeout: float = 5.0,
    poll_interval: float = 0.05,
) -> int | None:
    """Poll query_lookup_result until it returns a non-None value."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        result = ctrl.query_lookup_result(req_id)
        if result is not None:
            return result
        time.sleep(poll_interval)
    return None


def wait_for_prefetch_result_bitmap(
    ctrl: PrefetchController,
    req_id: int,
    timeout: float = 5.0,
    poll_interval: float = 0.05,
):
    """Poll query_prefetch_result, returning the raw retained Bitmap.

    Unlike :func:`wait_for_prefetch_result`, this keeps the full bitmap so a
    caller can inspect non-contiguous retained sets (e.g. SEGMENTED_PREFIX).
    """
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        result = ctrl.query_prefetch_result(req_id)
        if result is not None:
            return result
        time.sleep(poll_interval)
    return None


def make_adapter() -> MockL2Adapter:
    """Create a MockL2Adapter with fast bandwidth."""
    config = MockL2AdapterConfig(max_size_gb=0.01, mock_bandwidth_gb=10.0)
    return MockL2Adapter(config)


def make_descriptor(index: int) -> AdapterDescriptor:
    """Create an AdapterDescriptor for testing."""
    config = MockL2AdapterConfig(max_size_gb=0.01, mock_bandwidth_gb=10.0)
    return AdapterDescriptor(index=index, config=config)


def store_keys_in_l2(
    adapter: MockL2Adapter,
    keys: list[ObjectKey],
    layout: MemoryLayoutDesc,
) -> None:
    """Store test data directly in L2 adapter and wait for completion."""
    if not keys:
        return
    objs = []
    for _ in keys:
        tensor = torch.randn(layout.shapes[0], dtype=layout.dtypes[0])
        metadata = MemoryObjMetadata(
            shape=layout.shapes[0],
            dtype=layout.dtypes[0],
            address=0,
            phy_size=tensor.nelement() * tensor.element_size(),
            ref_count=0,
        )
        obj = TensorMemoryObj(raw_data=tensor, metadata=metadata, parent_allocator=None)
        objs.append(obj)
    adapter.submit_store_task(keys, objs)  # type: ignore
    ok = wait_for_condition(
        lambda: all(adapter.debug_has_key(k) for k in keys),
        timeout=5.0,
    )
    assert ok, "Failed to store test data in L2 adapter"


# =============================================================================
# Fixtures
# =============================================================================


@pytest.fixture
def l1_manager():
    """Create an L1Manager with a reasonable memory config."""
    config = L1ManagerConfig(
        memory_config=L1MemoryManagerConfig(
            size_in_bytes=128 * 1024 * 1024,
            use_lazy=should_use_lazy_alloc(),
            init_size_in_bytes=64 * 1024 * 1024,
            align_bytes=0x1000,
        ),
        write_ttl_seconds=600,
        read_ttl_seconds=300,
    )
    mgr = L1Manager(config)
    yield mgr
    mgr.close()


# =============================================================================
# Lifecycle Tests
# =============================================================================


class TestPrefetchControllerLifecycle:
    """Test PrefetchController start/stop behavior."""

    def test_start_stop(self, l1_manager):
        """Controller should start and stop cleanly."""
        adapter = make_adapter()
        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()
        ctrl.stop()
        adapter.close()

    def test_start_stop_no_adapters(self, l1_manager):
        """Controller should start and stop cleanly with no adapters."""
        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[],
            adapter_descriptors=[],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()
        ctrl.stop()


# =============================================================================
# Single Adapter Prefetch
# =============================================================================


class TestSingleAdapterPrefetch:
    """Test PrefetchController with one MockL2Adapter."""

    def test_full_prefix_hit(self, l1_manager):
        """All keys in L2 → all loaded, prefix hits = total keys."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(5)]

        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)

        assert result == 5, f"Expected 5 prefix hits, got {result}"

        # Verify prefix keys are read-locked in L1
        read_results = l1_manager.unsafe_read(keys)
        for key in keys:
            assert read_results[key][0] == L1Error.SUCCESS

        # Cleanup: release read locks
        l1_manager.finish_read(keys)

        ctrl.stop()
        adapter.close()

    def test_prefix_with_gap(self, l1_manager):
        """L2 has keys {0,1,3,4} but not 2 → only prefix {0,1} loaded."""
        adapter = make_adapter()
        layout = make_layout()
        all_keys = [make_object_key(i) for i in range(5)]
        # Store only keys 0, 1, 3, 4 (gap at index 2)
        stored_keys = [all_keys[i] for i in [0, 1, 3, 4]]
        store_keys_in_l2(adapter, stored_keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(all_keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)

        assert result == 2, f"Expected 2 prefix hits (gap at index 2), got {result}"

        # Verify prefix keys {0,1} are read-locked in L1
        prefix_keys = all_keys[:2]
        read_results = l1_manager.unsafe_read(prefix_keys)
        for key in prefix_keys:
            assert read_results[key][0] == L1Error.SUCCESS

        # Verify keys beyond prefix {2,3,4} are NOT in L1
        non_prefix_keys = all_keys[2:]
        read_results = l1_manager.reserve_read(non_prefix_keys)
        for key in non_prefix_keys:
            assert read_results[key][0] == L1Error.KEY_NOT_EXIST

        l1_manager.finish_read(prefix_keys)
        ctrl.stop()
        adapter.close()

    def test_segmented_prefix_with_gap(self, l1_manager):
        """SEGMENTED_PREFIX retains the post-gap keys: L2 {0,1,3,4} -> {0,1,3,4}.

        The PREFIX counterpart (test_prefix_with_gap) truncates the same gap to
        {0,1}; SEGMENTED_PREFIX keeps the post-gap chunks L1-resident so only
        the hole (index 2) needs recomputing.
        """
        adapter = make_adapter()
        layout = make_layout()
        all_keys = [make_object_key(i) for i in range(5)]
        # Store only keys 0, 1, 3, 4 (gap at index 2).
        store_keys_in_l2(adapter, [all_keys[i] for i in [0, 1, 3, 4]], layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(
            all_keys, layout, policy=TrimPolicy.SEGMENTED_PREFIX
        )
        retained = wait_for_prefetch_result_bitmap(ctrl, req_id)
        assert retained is not None
        assert retained.get_indices_list() == [0, 1, 3, 4], (
            "SEGMENTED_PREFIX should retain post-gap keys, got "
            f"{retained.get_indices_list()}"
        )

        # Retained keys {0,1,3,4} are read-locked in L1; the gap (2) is absent.
        retained_keys = [all_keys[i] for i in [0, 1, 3, 4]]
        read = l1_manager.unsafe_read(retained_keys)
        for key in retained_keys:
            assert read[key][0] == L1Error.SUCCESS
        gap_read = l1_manager.reserve_read([all_keys[2]])
        assert gap_read[all_keys[2]][0] == L1Error.KEY_NOT_EXIST

        l1_manager.finish_read(retained_keys)
        ctrl.stop()
        adapter.close()

    def test_fault_inject_load_gap_segmented_vs_prefix(self, l1_manager):
        """fault_inject (load fails at index 2) drives the segmented path.

        Lookup reports all 5 keys present; the *load* of index 2 fails (the L2
        retrieve error the adapter simulates). SEGMENTED_PREFIX retains the
        post-gap keys {0,1,3,4}; PREFIX truncates at the hole to {0,1}. Distinct
        key ranges per policy keep the shared L1 from serving the second pass.
        """
        layout = make_layout()
        for base, trim, expected in (
            (0, TrimPolicy.SEGMENTED_PREFIX, [0, 1, 3, 4]),
            (10, TrimPolicy.PREFIX, [0, 1]),
        ):
            keys = [make_object_key(base + i) for i in range(5)]
            inner = make_adapter()
            store_keys_in_l2(inner, keys, layout)  # all 5 present at lookup
            # Drop task-position 2 at load -> a mid-prefix L2 retrieve failure.
            fault = FaultInjectL2Adapter(inner, rate=0.0, seed=0, gap_indices=(2,))

            ctrl = PrefetchController(
                l1_manager=l1_manager,
                l2_adapters=[fault],
                adapter_descriptors=[make_descriptor(0)],
                policy=DefaultPrefetchPolicy(),
            )
            ctrl.start()
            req_id = ctrl.submit_prefetch_request(keys, layout, policy=trim)
            retained = wait_for_prefetch_result_bitmap(ctrl, req_id)
            assert retained is not None
            assert retained.get_indices_list() == expected, (
                f"{trim.name}: expected {expected}, got {retained.get_indices_list()}"
            )

            held = [keys[i] for i in retained.get_indices_list()]
            if held:
                l1_manager.finish_read(held)
            ctrl.stop()
            fault.close()

    def test_key0_missing(self, l1_manager):
        """L2 has keys {1,2,3} but not 0 → prefix = 0, nothing loaded."""
        adapter = make_adapter()
        layout = make_layout()
        all_keys = [make_object_key(i) for i in range(4)]
        # Store keys 1, 2, 3 but not 0
        stored_keys = all_keys[1:]
        store_keys_in_l2(adapter, stored_keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(all_keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)

        assert result == 0, f"Expected 0 prefix hits (key 0 missing), got {result}"

        # Verify no keys are in L1
        read_results = l1_manager.reserve_read(all_keys)
        for key in all_keys:
            assert read_results[key][0] == L1Error.KEY_NOT_EXIST

        ctrl.stop()
        adapter.close()


# =============================================================================
# Multi Adapter Prefetch
# =============================================================================


class TestMultiAdapterPrefetch:
    """Test PrefetchController with multiple MockL2Adapters."""

    def test_disjoint_adapters(self, l1_manager):
        """Adapter 0 has {0,1}, adapter 1 has {2,3} → full prefix of 4."""
        adapters = [make_adapter(), make_adapter()]
        descriptors = [make_descriptor(i) for i in range(2)]
        layout = make_layout()
        keys = [make_object_key(i) for i in range(4)]

        store_keys_in_l2(adapters[0], keys[:2], layout)
        store_keys_in_l2(adapters[1], keys[2:], layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=adapters,
            adapter_descriptors=descriptors,
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)

        assert result == 4, f"Expected 4 prefix hits, got {result}"

        read_results = l1_manager.unsafe_read(keys)
        for key in keys:
            assert read_results[key][0] == L1Error.SUCCESS

        l1_manager.finish_read(keys)
        ctrl.stop()
        for a in adapters:
            a.close()

    def test_overlap_first_wins(self, l1_manager):
        """Both adapters have key 1. Adapter 0 (lower index) loads it."""
        adapters = [make_adapter(), make_adapter()]
        descriptors = [make_descriptor(i) for i in range(2)]
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]

        # Adapter 0 has keys {0, 1}, adapter 1 has keys {1, 2}
        store_keys_in_l2(adapters[0], keys[:2], layout)
        store_keys_in_l2(adapters[1], keys[1:], layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=adapters,
            adapter_descriptors=descriptors,
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)

        assert result == 3, f"Expected 3 prefix hits, got {result}"

        read_results = l1_manager.unsafe_read(keys)
        for key in keys:
            assert read_results[key][0] == L1Error.SUCCESS

        l1_manager.finish_read(keys)
        ctrl.stop()
        for a in adapters:
            a.close()


# =============================================================================
# No Hits
# =============================================================================


class TestNoHits:
    """Test PrefetchController when no keys are found in L2."""

    def test_no_keys_in_l2(self, l1_manager):
        """Prefetch keys not in any L2 → 0 prefix hits."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]
        # Don't store anything in L2

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)

        assert result == 0, f"Expected 0 prefix hits, got {result}"

        ctrl.stop()
        adapter.close()

    def test_no_adapters(self, l1_manager):
        """No adapters → 0 prefix hits immediately."""
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[],
            adapter_descriptors=[],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)

        assert result == 0, f"Expected 0 prefix hits, got {result}"

        ctrl.stop()


# =============================================================================
# Query Result
# =============================================================================


class TestQueryResult:
    """Test query_prefetch_result semantics."""

    def test_query_returns_int_then_none(self, l1_manager):
        """Result is consumed on first query; second query returns None."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(2)]
        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)

        assert result == 2
        # Second query should return None (already consumed)
        assert ctrl.query_prefetch_result(req_id) is None

        l1_manager.finish_read(keys)
        ctrl.stop()
        adapter.close()

    def test_query_before_completion_returns_none(self, l1_manager):
        """Querying a nonexistent request ID returns None."""
        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[],
            adapter_descriptors=[],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        assert ctrl.query_prefetch_result(999) is None

        ctrl.stop()


# =============================================================================
# L2 Lock Release
# =============================================================================


class TestPrefetchL2LockRelease:
    """Test that L2 locks are properly released after prefetch."""

    def test_locks_released_after_full_hit(self, l1_manager):
        """All L2 locks should be released after a successful prefetch."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]
        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 3

        # Wait for L2 unlock operations to be processed
        ok = wait_for_condition(
            lambda: adapter.debug_get_locked_key_count() == 0,
            timeout=5.0,
        )
        assert ok, "L2 locks should be released after prefetch completion"

        l1_manager.finish_read(keys)
        ctrl.stop()
        adapter.close()

    def test_locks_released_after_prefix_trim(self, l1_manager):
        """L2 locks released for both prefix and non-prefix keys."""
        adapter = make_adapter()
        layout = make_layout()
        all_keys = [make_object_key(i) for i in range(5)]
        # Store keys 0, 1, 3, 4 (gap at index 2)
        stored_keys = [all_keys[i] for i in [0, 1, 3, 4]]
        store_keys_in_l2(adapter, stored_keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(all_keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 2

        # All L2 locks should be released (both prefix and trimmed keys)
        ok = wait_for_condition(
            lambda: adapter.debug_get_locked_key_count() == 0,
            timeout=5.0,
        )
        assert ok, "All L2 locks should be released after prefix-trimmed prefetch"

        l1_manager.finish_read(all_keys[:2])
        ctrl.stop()
        adapter.close()

    def test_locks_released_after_no_hits(self, l1_manager):
        """L2 locks released even when nothing is found."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]
        # Don't store anything

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 0

        ok = wait_for_condition(
            lambda: adapter.debug_get_locked_key_count() == 0,
            timeout=5.0,
        )
        assert ok, "L2 locks should be 0 when nothing was found"

        ctrl.stop()
        adapter.close()

    def test_multi_adapter_locks_released(self, l1_manager):
        """Both adapters' L2 locks released after overlapping prefetch."""
        adapters = [make_adapter(), make_adapter()]
        descriptors = [make_descriptor(i) for i in range(2)]
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]

        # Adapter 0 has {0, 1}, adapter 1 has {1, 2}
        store_keys_in_l2(adapters[0], keys[:2], layout)
        store_keys_in_l2(adapters[1], keys[1:], layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=adapters,
            adapter_descriptors=descriptors,
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 3

        # Both adapters should have all locks released
        for i, adapter in enumerate(adapters):
            ok = wait_for_condition(
                lambda a=adapter: a.debug_get_locked_key_count() == 0,
                timeout=5.0,
            )
            assert ok, f"Adapter {i} should have all L2 locks released"

        l1_manager.finish_read(keys)
        ctrl.stop()
        for a in adapters:
            a.close()


# =============================================================================
# Max In-Flight
# =============================================================================


class TestMaxInFlight:
    """Test PrefetchController max in-flight request limiting."""

    def test_queuing_beyond_max_in_flight(self, l1_manager):
        """Submit more requests than max_in_flight → all eventually complete."""
        adapter = make_adapter()
        layout = make_layout()

        # Store keys for 4 separate requests (2 keys each)
        all_keys = [make_object_key(i) for i in range(8)]
        store_keys_in_l2(adapter, all_keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
            max_in_flight=2,
        )
        ctrl.start()

        # Submit 4 requests (max_in_flight=2, so 2 queued)
        req_ids = []
        for i in range(4):
            batch_keys = all_keys[i * 2 : (i + 1) * 2]
            req_id = ctrl.submit_prefetch_request(batch_keys, layout)
            req_ids.append(req_id)

        # All 4 requests should eventually complete
        results = []
        for req_id in req_ids:
            result = wait_for_prefetch_result(ctrl, req_id, timeout=10.0)
            assert result is not None, f"Request {req_id} should complete"
            results.append(result)

        assert results == [2, 2, 2, 2]

        # Release read locks for all keys
        l1_manager.finish_read(all_keys)

        ctrl.stop()
        adapter.close()


# =============================================================================
# Multiple Sequential Requests
# =============================================================================


class TestMultipleRequests:
    """Test multiple sequential prefetch requests."""

    def test_two_sequential_requests(self, l1_manager):
        """Two back-to-back requests should both complete correctly."""
        adapter = make_adapter()
        layout = make_layout()
        keys1 = [make_object_key(i) for i in range(3)]
        keys2 = [make_object_key(i) for i in range(10, 14)]

        store_keys_in_l2(adapter, keys1, layout)
        store_keys_in_l2(adapter, keys2, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req1 = ctrl.submit_prefetch_request(keys1, layout)
        result1 = wait_for_prefetch_result(ctrl, req1)
        assert result1 == 3

        req2 = ctrl.submit_prefetch_request(keys2, layout)
        result2 = wait_for_prefetch_result(ctrl, req2)
        assert result2 == 4

        l1_manager.finish_read(keys1)
        l1_manager.finish_read(keys2)
        ctrl.stop()
        adapter.close()


# =============================================================================
# extra_count Path
# =============================================================================


class TestExtraCountPrefetch:
    """Test that extra_count is correctly propagated through the prefetch path.

    When extra_count=N is passed to submit_prefetch_request, the controller
    must acquire 1 + N read locks per key (one for the prefetch controller
    itself, plus N for additional TP workers).  Each consumer must call
    finish_read (or finish_read_prefetched) once to release its lock.

    These tests verify:
    1. Keys remain accessible after the first finish_read when extra_count > 0.
    2. Keys are evictable only after ALL 1 + N locks are released.
    3. extra_count=0 (default) behaves identically to the original single-lock
       path.
    4. Prefix trimming still works correctly with extra_count > 0.
    5. Non-prefix loaded keys have all extra locks released by _finalize_load.
    """

    def test_extra_count_zero_default_behavior(self, l1_manager):
        """extra_count=0 (default): single read lock, key freed after one
        finish_read."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]
        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout, extra_count=0)
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 3

        # Keys should be readable immediately after prefetch
        read_results = l1_manager.unsafe_read(keys)
        for key in keys:
            assert read_results[key][0] == L1Error.SUCCESS

        # Release the single read lock — keys should become unlocked
        finish_results = l1_manager.finish_read(keys, extra_count=0)
        for key in keys:
            assert finish_results[key] == L1Error.SUCCESS

        ctrl.stop()
        adapter.close()

    def test_extra_count_one_requires_two_finish_reads(self, l1_manager):
        """extra_count=1: two read locks acquired; key stays locked after
        first finish_read and is released after second."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]
        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        # extra_count=1 → 2 read locks per key
        req_id = ctrl.submit_prefetch_request(keys, layout, extra_count=1)
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 3

        # Keys must be readable right after prefetch
        read_results = l1_manager.unsafe_read(keys)
        for key in keys:
            assert read_results[key][0] == L1Error.SUCCESS

        # Release lock #1 (the "prefetch controller" lock, extra_count=0)
        finish_results = l1_manager.finish_read(keys, extra_count=0)
        for key in keys:
            assert finish_results[key] == L1Error.SUCCESS

        # Keys must STILL be readable — lock #2 (the TP worker lock) is held
        read_results2 = l1_manager.unsafe_read(keys)
        for key in keys:
            assert read_results2[key][0] == L1Error.SUCCESS, (
                f"Key {key} should still be read-locked after first finish_read"
            )

        # Release lock #2 (the TP worker lock, extra_count=0)
        finish_results2 = l1_manager.finish_read(keys, extra_count=0)
        for key in keys:
            assert finish_results2[key] == L1Error.SUCCESS

        ctrl.stop()
        adapter.close()

    def test_extra_count_three_requires_four_finish_reads(self, l1_manager):
        """extra_count=3 (TP=4): four read locks; key stays locked until all
        four are released."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(2)]
        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        # extra_count=3 → 4 read locks per key
        req_id = ctrl.submit_prefetch_request(keys, layout, extra_count=3)
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 2

        # Release locks one by one; key must remain readable until the last
        for release_idx in range(3):
            finish_results = l1_manager.finish_read(keys, extra_count=0)
            for key in keys:
                assert finish_results[key] == L1Error.SUCCESS

            # Still readable — remaining locks are held
            read_results = l1_manager.unsafe_read(keys)
            for key in keys:
                assert read_results[key][0] == L1Error.SUCCESS, (
                    f"Key {key} should still be locked after {release_idx + 1} "
                    f"finish_read calls"
                )

        # Release the final lock
        finish_results = l1_manager.finish_read(keys, extra_count=0)
        for key in keys:
            assert finish_results[key] == L1Error.SUCCESS

        ctrl.stop()
        adapter.close()

    def test_extra_count_with_prefix_trim(self, l1_manager):
        """extra_count=1 with a gap in L2: only prefix keys get 2 locks;
        non-prefix keys are never loaded."""
        adapter = make_adapter()
        layout = make_layout()
        all_keys = [make_object_key(i) for i in range(5)]
        # Store keys 0, 1, 3, 4 — gap at index 2
        stored_keys = [all_keys[i] for i in [0, 1, 3, 4]]
        store_keys_in_l2(adapter, stored_keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(all_keys, layout, extra_count=1)
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 2, f"Expected 2 prefix hits (gap at index 2), got {result}"

        prefix_keys = all_keys[:2]

        # Prefix keys must be readable
        read_results = l1_manager.unsafe_read(prefix_keys)
        for key in prefix_keys:
            assert read_results[key][0] == L1Error.SUCCESS

        # Release lock #1 — prefix keys still held by lock #2
        l1_manager.finish_read(prefix_keys, extra_count=0)

        read_results2 = l1_manager.unsafe_read(prefix_keys)
        for key in prefix_keys:
            assert read_results2[key][0] == L1Error.SUCCESS, (
                f"Prefix key {key} should still be locked after first finish_read"
            )

        # Release lock #2
        l1_manager.finish_read(prefix_keys, extra_count=0)

        # Non-prefix keys must NOT be in L1
        non_prefix_keys = all_keys[2:]
        reserve_results = l1_manager.reserve_read(non_prefix_keys)
        for key in non_prefix_keys:
            assert reserve_results[key][0] == L1Error.KEY_NOT_EXIST, (
                f"Non-prefix key {key} should not be in L1"
            )

        ctrl.stop()
        adapter.close()

    def test_extra_count_non_prefix_loaded_keys_fully_released(self, l1_manager):
        """Keys loaded beyond the prefix (due to partial load failure) must
        have ALL extra locks released by _finalize_load so they can be evicted.

        We simulate this by storing keys {0, 1, 2} in L2 but making key 1
        fail to reserve in L1 (by pre-occupying it), creating a gap so that
        key 2 is loaded but lies beyond the prefix.  _finalize_load must
        release 1 + extra_count locks for key 2.
        """
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]
        store_keys_in_l2(adapter, keys, layout)

        # Pre-occupy key[1] in L1 with a write lock so reserve_write fails for it.
        # This forces a gap: key 0 is prefix (hit), key 1 fails reservation
        # (gap), key 2 is loaded but beyond the prefix.
        pre_write_results = l1_manager.reserve_write(
            keys=[keys[1]],
            is_temporary=[False],
            layout_desc=layout,
            mode="new",
        )
        assert pre_write_results[keys[1]][0] == L1Error.SUCCESS

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout, extra_count=1)
        result = wait_for_prefetch_result(ctrl, req_id)
        # Only key 0 is in the prefix (key 1 reservation failed → gap)
        assert result == 1, f"Expected 1 prefix hit, got {result}"

        # key[0] should be in L1 with 2 read locks (1 + extra_count=1)
        read_results = l1_manager.unsafe_read([keys[0]])
        assert read_results[keys[0]][0] == L1Error.SUCCESS

        # key[2] was loaded but is beyond the prefix; _finalize_load must have
        # released all 1 + extra_count=2 locks, so it should be gone from L1
        # (it's a temporary object and its lock count should be 0).
        reserve_results = l1_manager.reserve_read([keys[2]])
        assert reserve_results[keys[2]][0] == L1Error.KEY_NOT_EXIST, (
            "Non-prefix loaded key[2] should have all locks released and be "
            "evicted from L1"
        )

        # Clean up: release key[0]'s 2 locks and key[1]'s write lock
        l1_manager.finish_read([keys[0]], extra_count=0)
        l1_manager.finish_read([keys[0]], extra_count=0)
        l1_manager.finish_write([keys[1]])
        l1_manager.delete([keys[1]])

        ctrl.stop()
        adapter.close()


# =============================================================================
# Query Lookup Result
# =============================================================================


class TestQueryLookupResult:
    """Test query_lookup_result semantics."""

    def test_lookup_result_available_before_prefetch_completes(self, l1_manager):
        """Lookup result is available while load is still in progress."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]
        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)

        # Lookup result should be available before or at the same time as
        # the full prefetch result.
        lookup_hits = wait_for_lookup_result(ctrl, req_id)
        assert lookup_hits is not None
        assert lookup_hits == 3

        # Wait for full prefetch to complete and clean up
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 3

        l1_manager.finish_read(keys)
        ctrl.stop()
        adapter.close()

    def test_lookup_result_with_prefix_gap(self, l1_manager):
        """Lookup result reflects prefix-only hits (gap breaks prefix)."""
        adapter = make_adapter()
        layout = make_layout()
        # Store keys 0 and 2 (gap at 1)
        keys = [make_object_key(i) for i in range(3)]
        store_keys_in_l2(adapter, [keys[0], keys[2]], layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        lookup_hits = wait_for_lookup_result(ctrl, req_id)
        assert lookup_hits is not None
        # Only key 0 is in the prefix (gap at key 1 breaks it)
        assert lookup_hits == 1

        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 1

        l1_manager.finish_read([keys[0]])
        ctrl.stop()
        adapter.close()

    def test_lookup_result_zero_hits(self, l1_manager):
        """Lookup result is 0 when no keys are found in L2."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]
        # Don't store anything in L2

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        lookup_hits = wait_for_lookup_result(ctrl, req_id)
        assert lookup_hits is not None
        assert lookup_hits == 0

        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 0

        ctrl.stop()
        adapter.close()

    def test_lookup_result_not_popped_by_query(self, l1_manager):
        """query_lookup_result does not consume the result (idempotent reads)."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(2)]
        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        lookup_hits = wait_for_lookup_result(ctrl, req_id)
        assert lookup_hits == 2

        # Second query should still return the same value (not consumed)
        assert ctrl.query_lookup_result(req_id) == 2

        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 2

        l1_manager.finish_read(keys)
        ctrl.stop()
        adapter.close()

    def test_lookup_result_cleaned_up_by_prefetch_result(self, l1_manager):
        """query_prefetch_result cleans up the lookup result entry."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(2)]
        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        lookup_hits = wait_for_lookup_result(ctrl, req_id)
        assert lookup_hits == 2

        # Consume the prefetch result (should also clean up lookup entry)
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 2

        # Lookup result should now be gone
        assert ctrl.query_lookup_result(req_id) is None

        l1_manager.finish_read(keys)
        ctrl.stop()
        adapter.close()

    def test_lookup_result_nonexistent_request(self, l1_manager):
        """Querying a nonexistent request ID returns None."""
        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[],
            adapter_descriptors=[],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        assert ctrl.query_lookup_result(999) is None

        ctrl.stop()


class TestBuildTrimMask:
    """build_trim_mask picks the retained subset per policy: PREFIX trims at
    the first gap; SEGMENTED_PREFIX and SPARSE keep every set bit (gaps and
    all). The retained bitmap is consumed unchanged at the controller's load
    sites, so testing the mask directly covers the policy semantics."""

    @staticmethod
    def _bm(n, idxs):
        bm = Bitmap(n)
        for i in idxs:
            bm.set(i)
        return bm

    def test_prefix_trims_at_first_gap(self):
        found = self._bm(5, [0, 1, 3, 4])  # gap at index 2
        assert build_trim_mask(found, 5, TrimPolicy.PREFIX).get_indices_list() == [
            0,
            1,
        ]

    def test_segmented_prefix_keeps_gaps(self):
        # Models an L2 hit whose L1 load failed mid-prefix (e.g. OOM at index
        # 2): the keys that did load are kept, not trimmed to the first gap.
        found = self._bm(5, [0, 1, 3, 4])
        assert build_trim_mask(
            found, 5, TrimPolicy.SEGMENTED_PREFIX
        ).get_indices_list() == [0, 1, 3, 4]

    def test_sparse_keeps_all_found(self):
        found = self._bm(5, [0, 2, 4])
        assert build_trim_mask(found, 5, TrimPolicy.SPARSE).get_indices_list() == [
            0,
            2,
            4,
        ]


class TestMergeBitmaps:
    """merge_bitmaps always returns a num_keys-sized bitmap."""

    def test_empty_input_returns_sized_bitmap(self):
        """Empty input -> num_keys-sized all-zeros bitmap (not Bitmap(0)), so a
        downstream ``&`` with a same-sized mask never hits a size mismatch."""
        merged = merge_bitmaps([], 5)
        assert merged.popcount() == 0
        mask = Bitmap(5)
        mask.set(2)
        assert (merged & mask).popcount() == 0  # would raise on size mismatch

    def test_empty_generator_returns_sized_bitmap(self):
        """A generator is truthy even when empty; the result is still size-5."""
        merged = merge_bitmaps((b for b in []), 5)
        assert merged.popcount() == 0
        assert (merged & Bitmap(5)).popcount() == 0

    def test_union_of_bitmaps(self):
        """Non-empty inputs are OR-merged into one num_keys-sized bitmap."""
        a, b = Bitmap(5), Bitmap(5)
        a.set(0)
        b.set(3)
        assert merge_bitmaps([a, b], 5).get_indices_list() == [0, 3]


class TestWaitPrefetchResult:
    """Test the blocking wait_prefetch_result interface."""

    def test_wait_blocks_until_result_ready(self, l1_manager):
        """wait_prefetch_result blocks until the background result is published,
        returns True, and does not consume the result."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(5)]
        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        # Blocks until the background thread publishes the result.
        assert ctrl.wait_prefetch_result(req_id, timeout=10.0) is True
        # wait_prefetch_result must not consume the result.
        result = ctrl.query_prefetch_result(req_id)
        assert result is not None
        assert result.count_leading_ones() == 5

        l1_manager.finish_read(keys)
        ctrl.stop()
        adapter.close()

    def test_wait_times_out_for_unknown_request(self, l1_manager):
        """wait_prefetch_result returns False, after genuinely waiting, when no
        result arrives within the timeout."""
        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[],
            adapter_descriptors=[],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        start = time.monotonic()
        assert ctrl.wait_prefetch_result(999999, timeout=0.2) is False
        assert time.monotonic() - start >= 0.2

        ctrl.stop()


# =============================================================================
# Retention Policy
# =============================================================================


class TestPrefetchMode:
    """``mode=WARM`` (the warm path) loads keys **permanent** and
    **without a read lock**, vs ``LOOKUP`` which read-locks temporary
    objects that vanish on release.

    Both tests use ``DefaultPrefetchPolicy`` so the only difference is the
    per-request ``mode`` argument.
    """

    def test_warm_loads_unlocked_and_permanent(self, l1_manager):
        """WARM loads keys permanent and with NO read lock: immediately ready
        (reserve_read SUCCEEDS), holding no lock (unsafe_read NOT_READABLE), and
        not deleted on a reserve_read/finish_read cycle (permanent, not temp)."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]
        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout, mode=PrefetchMode.WARM)
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 3

        # No warming lock: unsafe_read (which needs an active read lock) reports
        # NOT_READABLE even though the keys are present and ready.
        unsafe = l1_manager.unsafe_read(keys)
        for key in keys:
            assert unsafe[key][0] == L1Error.KEY_NOT_READABLE

        # They are ready and re-lookupable: reserve_read SUCCEEDS...
        read_results = l1_manager.reserve_read(keys)
        for key in keys:
            assert read_results[key][0] == L1Error.SUCCESS

        # ...and releasing that probe lock does NOT delete them (permanent).
        l1_manager.finish_read(keys)
        again = l1_manager.reserve_read(keys)
        for key in keys:
            assert again[key][0] == L1Error.SUCCESS

        l1_manager.finish_read(keys)
        l1_manager.delete(keys)
        ctrl.stop()
        adapter.close()

    def test_warm_sparse_skips_write_locked_prefix_and_loads_later_keys(
        self, l1_manager
    ):
        """SPARSE WARM does not let a write-locked earlier key suppress later
        keys that can be reserved and loaded."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]
        store_keys_in_l2(adapter, keys, layout)

        existing = l1_manager.reserve_write(
            [keys[0]], is_temporary=[False], layout_desc=layout, mode="new"
        )
        assert existing[keys[0]][0] == L1Error.SUCCESS

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(
            keys,
            layout,
            policy=TrimPolicy.SPARSE,
            mode=PrefetchMode.WARM,
        )
        result = wait_for_prefetch_result_bitmap(ctrl, req_id)
        assert result is not None
        assert result.get_indices_list() == [1, 2]

        l1_manager.finish_write([keys[0]])
        read_results = l1_manager.reserve_read(keys[1:])
        for key in keys[1:]:
            assert read_results[key][0] == L1Error.SUCCESS

        l1_manager.finish_read(keys[1:])
        l1_manager.delete(keys)
        ctrl.stop()
        adapter.close()

    def test_default_deletes_keys_after_finish_read(self, l1_manager):
        """LOOKUP defers to ``DefaultPrefetchPolicy`` (temporary), so
        the keys are deleted from L1 once the read-lock is released."""
        adapter = make_adapter()
        layout = make_layout()
        keys = [make_object_key(i) for i in range(3)]
        store_keys_in_l2(adapter, keys, layout)

        ctrl = PrefetchController(
            l1_manager=l1_manager,
            l2_adapters=[adapter],
            adapter_descriptors=[make_descriptor(0)],
            policy=DefaultPrefetchPolicy(),
        )
        ctrl.start()

        req_id = ctrl.submit_prefetch_request(keys, layout)
        result = wait_for_prefetch_result(ctrl, req_id)
        assert result == 3

        # Releasing the read-lock deletes the temporary objects, so a
        # subsequent lookup (reserve_read) misses with KEY_NOT_EXIST.
        l1_manager.finish_read(keys)
        read_results = l1_manager.reserve_read(keys)
        for key in keys:
            assert read_results[key][0] == L1Error.KEY_NOT_EXIST

        ctrl.stop()
        adapter.close()