chia-blockchain/chia/_tests/plot_sync/test_sync_simulated.py

from __future__ import annotations

import asyncio
import contextlib
import functools
import logging
import random
import time
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any

import pytest
from chia_rs import G1Element
from chia_rs.sized_bytes import bytes32
from chia_rs.sized_ints import int16, uint8, uint64

from chia._tests.plot_sync.util import start_harvester_service
from chia._tests.util.time_out_assert import time_out_assert
from chia.farmer.farmer import Farmer
from chia.farmer.farmer_service import FarmerService
from chia.harvester.harvester import Harvester
from chia.harvester.harvester_service import HarvesterService
from chia.plot_sync.receiver import Receiver
from chia.plot_sync.sender import Sender
from chia.plot_sync.util import Constants
from chia.plotting.manager import PlotManager
from chia.plotting.prover import V1Prover
from chia.plotting.util import PlotInfo
from chia.protocols.harvester_protocol import PlotSyncError, PlotSyncResponse
from chia.protocols.outbound_message import make_msg
from chia.protocols.protocol_message_types import ProtocolMessageTypes
from chia.server.ws_connection import WSChiaConnection
from chia.simulator.block_tools import BlockTools
from chia.util.batches import to_batches

log = logging.getLogger(__name__)


class ErrorSimulation(Enum):
    DropEveryFourthMessage = 1
    DropThreeMessages = 2
    RespondTooLateEveryFourthMessage = 3
    RespondTwice = 4
    NonRecoverableError = 5
    NotConnected = 6


@dataclass
class TestData:
    harvester: Harvester
    plot_sync_sender: Sender
    plot_sync_receiver: Receiver
    event_loop: asyncio.AbstractEventLoop
    plots: dict[Path, PlotInfo] = field(default_factory=dict)
    invalid: list[PlotInfo] = field(default_factory=list)
    keys_missing: list[PlotInfo] = field(default_factory=list)
    duplicates: list[PlotInfo] = field(default_factory=list)

    async def run(
        self,
        *,
        loaded: list[PlotInfo],
        removed: list[PlotInfo],
        invalid: list[PlotInfo],
        keys_missing: list[PlotInfo],
        duplicates: list[PlotInfo],
        initial: bool,
    ) -> None:
        for plot_info in loaded:
            assert Path(plot_info.prover.get_filename()) not in self.plots
        for plot_info in removed:
            assert Path(plot_info.prover.get_filename()) in self.plots

        self.invalid = invalid
        self.keys_missing = keys_missing
        self.duplicates = duplicates

        removed_paths: list[Path] = [Path(p.prover.get_filename()) for p in removed] if removed is not None else []
        invalid_dict: dict[Path, int] = {Path(p.prover.get_filename()): 0 for p in self.invalid}
        keys_missing_set: set[Path] = {Path(p.prover.get_filename()) for p in self.keys_missing}
        duplicates_set: set[str] = {p.prover.get_filename() for p in self.duplicates}

        # Inject invalid plots into `PlotManager` of the harvester so that the callback calls below can use them
        # to sync them to the farmer.
        self.harvester.plot_manager.failed_to_open_filenames = invalid_dict
        # Inject key missing plots into `PlotManager` of the harvester so that the callback calls below can use them
        # to sync them to the farmer.
        self.harvester.plot_manager.no_key_filenames = keys_missing_set
        # Inject duplicated plots into `PlotManager` of the harvester so that the callback calls below can use them
        # to sync them to the farmer.
        for plot_info in loaded:
            plot_path = Path(plot_info.prover.get_filename())
            self.harvester.plot_manager.plot_filename_paths[plot_path.name] = (str(plot_path.parent), set())
        for duplicate in duplicates_set:
            plot_path = Path(duplicate)
            assert plot_path.name in self.harvester.plot_manager.plot_filename_paths
            self.harvester.plot_manager.plot_filename_paths[plot_path.name][1].add(str(plot_path.parent))

        batch_size = self.harvester.plot_manager.refresh_parameter.batch_size

        # Used to capture the sync id in `run_internal`
        sync_id: uint64 | None = None

        def run_internal() -> None:
            nonlocal sync_id
            # Simulate one plot manager refresh cycle by calling the methods directly.
            self.harvester.plot_sync_sender.sync_start(len(loaded), initial)
            sync_id = self.plot_sync_sender._sync_id
            if len(loaded) == 0:
                self.harvester.plot_sync_sender.process_batch([], 0)
            for batch in to_batches(loaded, batch_size):
                self.harvester.plot_sync_sender.process_batch(batch.entries, batch.remaining)
            self.harvester.plot_sync_sender.sync_done(removed_paths, 0)

        await self.event_loop.run_in_executor(None, run_internal)

        async def sync_done() -> bool:
            assert sync_id is not None
            return self.plot_sync_receiver.last_sync().sync_id == self.plot_sync_sender._last_sync_id == sync_id

        await time_out_assert(60, sync_done)

        for plot_info in loaded:
            self.plots[Path(plot_info.prover.get_filename())] = plot_info
        for plot_info in removed:
            del self.plots[Path(plot_info.prover.get_filename())]

    def validate_plot_sync(self) -> None:
        assert len(self.plots) == len(self.plot_sync_receiver.plots())
        assert len(self.invalid) == len(self.plot_sync_receiver.invalid())
        assert len(self.keys_missing) == len(self.plot_sync_receiver.keys_missing())
        for _, plot_info in self.plots.items():
            assert plot_info.prover.get_filename() not in self.plot_sync_receiver.invalid()
            assert plot_info.prover.get_filename() not in self.plot_sync_receiver.keys_missing()
            assert plot_info.prover.get_filename() in self.plot_sync_receiver.plots()
            synced_plot = self.plot_sync_receiver.plots()[plot_info.prover.get_filename()]
            assert plot_info.prover.get_filename() == synced_plot.filename
            assert plot_info.pool_public_key == synced_plot.pool_public_key
            assert plot_info.pool_contract_puzzle_hash == synced_plot.pool_contract_puzzle_hash
            assert plot_info.plot_public_key == synced_plot.plot_public_key
            assert plot_info.file_size == synced_plot.file_size
            assert uint64(plot_info.time_modified) == synced_plot.time_modified
        for plot_info in self.invalid:
            assert plot_info.prover.get_filename() not in self.plot_sync_receiver.plots()
            assert plot_info.prover.get_filename() in self.plot_sync_receiver.invalid()
            assert plot_info.prover.get_filename() not in self.plot_sync_receiver.keys_missing()
            assert plot_info.prover.get_filename() not in self.plot_sync_receiver.duplicates()
        for plot_info in self.keys_missing:
            assert plot_info.prover.get_filename() not in self.plot_sync_receiver.plots()
            assert plot_info.prover.get_filename() not in self.plot_sync_receiver.invalid()
            assert plot_info.prover.get_filename() in self.plot_sync_receiver.keys_missing()
            assert plot_info.prover.get_filename() not in self.plot_sync_receiver.duplicates()
        for plot_info in self.duplicates:
            assert plot_info.prover.get_filename() not in self.plot_sync_receiver.invalid()
            assert plot_info.prover.get_filename() not in self.plot_sync_receiver.keys_missing()
            assert plot_info.prover.get_filename() in self.plot_sync_receiver.duplicates()


@dataclass
class TestRunner:
    test_data: list[TestData]

    def __init__(
        self, harvesters: list[Harvester], farmer: Farmer, event_loop: asyncio.events.AbstractEventLoop
    ) -> None:
        self.test_data = []
        for harvester in harvesters:
            assert harvester.server is not None
            self.test_data.append(
                TestData(
                    harvester,
                    harvester.plot_sync_sender,
                    farmer.plot_sync_receivers[harvester.server.node_id],
                    event_loop,
                )
            )

    async def run(
        self,
        index: int,
        *,
        loaded: list[PlotInfo],
        removed: list[PlotInfo],
        invalid: list[PlotInfo],
        keys_missing: list[PlotInfo],
        duplicates: list[PlotInfo],
        initial: bool,
    ) -> None:
        await self.test_data[index].run(
            loaded=loaded,
            removed=removed,
            invalid=invalid,
            keys_missing=keys_missing,
            duplicates=duplicates,
            initial=initial,
        )
        for data in self.test_data:
            data.validate_plot_sync()


async def skip_processing(self: Any, _: WSChiaConnection, message_type: ProtocolMessageTypes, message: Any) -> bool:
    self.message_counter += 1
    if self.simulate_error == ErrorSimulation.DropEveryFourthMessage:
        if self.message_counter % 4 == 0:
            return True
    if self.simulate_error == ErrorSimulation.DropThreeMessages:
        if 2 < self.message_counter < 6:
            return True
    if self.simulate_error == ErrorSimulation.RespondTooLateEveryFourthMessage:
        if self.message_counter % 4 == 0:
            await asyncio.sleep(Constants.message_timeout + 1)
            return False
    if self.simulate_error == ErrorSimulation.RespondTwice:
        await self.connection().send_message(
            make_msg(
                ProtocolMessageTypes.plot_sync_response,
                PlotSyncResponse(message.identifier, int16(message_type.value), None),
            )
        )
    if self.simulate_error == ErrorSimulation.NonRecoverableError and self.message_counter > 1:
        await self.connection().send_message(
            make_msg(
                ProtocolMessageTypes.plot_sync_response,
                PlotSyncResponse(
                    message.identifier, int16(message_type.value), PlotSyncError(int16(0), "non recoverable", None)
                ),
            )
        )
        self.simulate_error = 0
        return True
    return False


async def _testable_process(
    self: Any, peer: WSChiaConnection, message_type: ProtocolMessageTypes, message: Any
) -> None:
    if await skip_processing(self, peer, message_type, message):
        return
    await self.original_process(peer, message_type, message)


@contextlib.asynccontextmanager
async def create_test_runner(
    harvester_services: list[HarvesterService],
    farmer_service: FarmerService,
    event_loop: asyncio.events.AbstractEventLoop,
) -> AsyncIterator[TestRunner]:
    async with farmer_service.manage():
        farmer: Farmer = farmer_service._node
        assert len(farmer.plot_sync_receivers) == 0
        async with contextlib.AsyncExitStack() as async_exit_stack:
            split_harvester_managers = [
                await async_exit_stack.enter_async_context(start_harvester_service(service, farmer_service))
                for service in harvester_services
            ]
            harvesters = [manager.object for manager in split_harvester_managers]
            for receiver in farmer.plot_sync_receivers.values():
                receiver.simulate_error = 0  # type: ignore[attr-defined]
                receiver.message_counter = 0  # type: ignore[attr-defined]
                receiver.original_process = receiver._process  # type: ignore[attr-defined]
                receiver._process = functools.partial(_testable_process, receiver)  # type: ignore[method-assign]
            yield TestRunner(harvesters, farmer, event_loop)


def create_example_plots(count: int, seeded_random: random.Random) -> list[PlotInfo]:
    @dataclass
    class DiskProver:
        file_name: str
        plot_id: bytes32
        size: int

        def get_filename(self) -> str:
            return self.file_name

        def get_id(self) -> bytes32:
            return self.plot_id

        def get_size(self) -> int:
            return self.size

        def get_compression_level(self) -> uint8:
            return uint8(0)

    return [
        PlotInfo(
            prover=V1Prover(DiskProver(f"{x}", bytes32.random(seeded_random), 25 + x % 26)),
            pool_public_key=None,
            pool_contract_puzzle_hash=None,
            plot_public_key=G1Element(),
            file_size=uint64(0),
            time_modified=time.time(),
        )
        for x in range(count)
    ]


@pytest.mark.anyio
async def test_sync_simulated(
    farmer_three_harvester_not_started: tuple[list[HarvesterService], FarmerService, BlockTools],
    event_loop: asyncio.events.AbstractEventLoop,
    seeded_random: random.Random,
) -> None:
    harvester_services, farmer_service, _ = farmer_three_harvester_not_started
    farmer: Farmer = farmer_service._node
    async with create_test_runner(harvester_services, farmer_service, event_loop) as test_runner:
        plots = create_example_plots(31000, seeded_random=seeded_random)

        await test_runner.run(
            0, loaded=plots[0:10000], removed=[], invalid=[], keys_missing=[], duplicates=plots[0:1000], initial=True
        )
        await test_runner.run(
            1,
            loaded=plots[10000:20000],
            removed=[],
            invalid=plots[30000:30100],
            keys_missing=[],
            duplicates=[],
            initial=True,
        )
        await test_runner.run(
            2,
            loaded=plots[20000:30000],
            removed=[],
            invalid=[],
            keys_missing=plots[30100:30200],
            duplicates=[],
            initial=True,
        )
        await test_runner.run(
            0,
            loaded=[],
            removed=[],
            invalid=plots[30300:30400],
            keys_missing=plots[30400:30453],
            duplicates=[],
            initial=False,
        )
        await test_runner.run(0, loaded=[], removed=[], invalid=[], keys_missing=[], duplicates=[], initial=False)
        await test_runner.run(
            0, loaded=[], removed=plots[5000:10000], invalid=[], keys_missing=[], duplicates=[], initial=False
        )
        await test_runner.run(
            1, loaded=[], removed=plots[10000:20000], invalid=[], keys_missing=[], duplicates=[], initial=False
        )
        await test_runner.run(
            2, loaded=[], removed=plots[20000:29000], invalid=[], keys_missing=[], duplicates=[], initial=False
        )
        await test_runner.run(
            0, loaded=[], removed=plots[0:5000], invalid=[], keys_missing=[], duplicates=[], initial=False
        )
        await test_runner.run(
            2,
            loaded=plots[5000:10000],
            removed=plots[29000:30000],
            invalid=plots[30000:30500],
            keys_missing=plots[30500:31000],
            duplicates=plots[5000:6000],
            initial=False,
        )
        await test_runner.run(
            2, loaded=[], removed=plots[5000:10000], invalid=[], keys_missing=[], duplicates=[], initial=False
        )
        assert len(farmer.plot_sync_receivers) == 3
        for plot_sync in farmer.plot_sync_receivers.values():
            assert len(plot_sync.plots()) == 0


@pytest.mark.parametrize(
    "simulate_error",
    [
        ErrorSimulation.DropEveryFourthMessage,
        ErrorSimulation.DropThreeMessages,
        ErrorSimulation.RespondTooLateEveryFourthMessage,
        ErrorSimulation.RespondTwice,
    ],
)
@pytest.mark.anyio
async def test_farmer_error_simulation(
    farmer_one_harvester_not_started: tuple[list[HarvesterService], FarmerService, BlockTools],
    event_loop: asyncio.events.AbstractEventLoop,
    simulate_error: ErrorSimulation,
    seeded_random: random.Random,
) -> None:
    Constants.message_timeout = 5
    harvester_services, farmer_service, _ = farmer_one_harvester_not_started
    async with create_test_runner(harvester_services, farmer_service, event_loop) as test_runner:
        batch_size = test_runner.test_data[0].harvester.plot_manager.refresh_parameter.batch_size
        plots = create_example_plots(batch_size + 3, seeded_random=seeded_random)
        receiver = test_runner.test_data[0].plot_sync_receiver
        receiver.simulate_error = simulate_error  # type: ignore[attr-defined]
        await test_runner.run(
            0,
            loaded=plots[0 : batch_size + 1],
            removed=[],
            invalid=[plots[batch_size + 1]],
            keys_missing=[plots[batch_size + 2]],
            duplicates=[],
            initial=True,
        )


@pytest.mark.parametrize("simulate_error", [ErrorSimulation.NonRecoverableError, ErrorSimulation.NotConnected])
@pytest.mark.anyio
async def test_sync_reset_cases(
    farmer_one_harvester_not_started: tuple[list[HarvesterService], FarmerService, BlockTools],
    event_loop: asyncio.events.AbstractEventLoop,
    simulate_error: ErrorSimulation,
    seeded_random: random.Random,
) -> None:
    harvester_services, farmer_service, _ = farmer_one_harvester_not_started
    async with create_test_runner(harvester_services, farmer_service, event_loop) as test_runner:
        test_data: TestData = test_runner.test_data[0]
        plot_manager: PlotManager = test_data.harvester.plot_manager
        plots = create_example_plots(30, seeded_random=seeded_random)
        # Inject some data into `PlotManager` of the harvester so that we can validate the reset worked and triggered a
        # fresh sync of all available data of the plot manager
        for plot_info in plots[0:10]:
            test_data.plots[Path(plot_info.prover.get_filename())] = plot_info
            plot_manager.plots = test_data.plots
        test_data.invalid = plots[10:20]
        test_data.keys_missing = plots[20:30]
        test_data.plot_sync_receiver.simulate_error = simulate_error  # type: ignore[attr-defined]
        sender: Sender = test_runner.test_data[0].plot_sync_sender
        started_sync_id: uint64 = uint64(0)

        plot_manager.failed_to_open_filenames = {Path(p.prover.get_filename()): 0 for p in test_data.invalid}
        plot_manager.no_key_filenames = {Path(p.prover.get_filename()) for p in test_data.keys_missing}

        async def wait_for_reset() -> bool:
            assert started_sync_id != 0
            return sender._sync_id != started_sync_id != 0

        async def sync_done() -> bool:
            assert started_sync_id != 0
            return test_data.plot_sync_receiver.last_sync().sync_id == sender._last_sync_id == started_sync_id

        # Send start and capture the sync_id
        sender.sync_start(len(plots), True)
        started_sync_id = sender._sync_id
        # Sleep 2 seconds to make sure we have a different sync_id after the reset which gets triggered
        await asyncio.sleep(2)
        saved_connection = sender._connection
        if simulate_error == ErrorSimulation.NotConnected:
            sender._connection = None
        sender.process_batch(plots, 0)
        await time_out_assert(60, wait_for_reset)
        started_sync_id = sender._sync_id
        sender._connection = saved_connection
        await time_out_assert(60, sync_done)
        test_runner.test_data[0].validate_plot_sync()