chia-blockchain/chia/plotting/cache.py

from __future__ import annotations

import logging
import time
import traceback
from collections.abc import ItemsView, KeysView, ValuesView
from dataclasses import dataclass, field
from functools import lru_cache
from math import ceil
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from chia.plotting.prover import ProverProtocol

from chia_rs import G1Element, PrivateKey
from chia_rs.sized_bytes import bytes32
from chia_rs.sized_ints import uint16, uint64

from chia.plotting.prover import get_prover_from_bytes
from chia.plotting.util import parse_plot_info
from chia.types.blockchain_format.proof_of_space import generate_plot_public_key
from chia.util.streamable import Streamable, VersionedBlob, streamable
from chia.wallet.derive_keys import master_sk_to_local_sk

log = logging.getLogger(__name__)

CURRENT_VERSION: int = 2


@lru_cache
def cached_master_sk_to_local_sk(master: PrivateKey) -> PrivateKey:
    return master_sk_to_local_sk(master)


@streamable
@dataclass(frozen=True)
class DiskCacheEntry(Streamable):
    prover_data: bytes
    farmer_public_key: G1Element
    pool_public_key: G1Element | None
    pool_contract_puzzle_hash: bytes32 | None
    plot_public_key: G1Element
    last_use: uint64


@streamable
@dataclass(frozen=True)
class CacheDataV1(Streamable):
    entries: list[tuple[str, DiskCacheEntry]]


@dataclass
class CacheEntry:
    prover: ProverProtocol
    farmer_public_key: G1Element
    pool_public_key: G1Element | None
    pool_contract_puzzle_hash: bytes32 | None
    plot_public_key: G1Element
    last_use: float

    @classmethod
    def from_prover(cls, prover: ProverProtocol) -> CacheEntry:
        (
            pool_public_key_or_puzzle_hash,
            farmer_public_key,
            local_master_sk,
        ) = parse_plot_info(prover.get_memo())

        pool_public_key: G1Element | None = None
        pool_contract_puzzle_hash: bytes32 | None = None
        if isinstance(pool_public_key_or_puzzle_hash, G1Element):
            pool_public_key = pool_public_key_or_puzzle_hash
        else:
            assert isinstance(pool_public_key_or_puzzle_hash, bytes32)
            pool_contract_puzzle_hash = pool_public_key_or_puzzle_hash

        local_sk = cached_master_sk_to_local_sk(local_master_sk)

        plot_public_key: G1Element = generate_plot_public_key(
            local_sk.get_g1(), farmer_public_key, pool_contract_puzzle_hash is not None
        )

        return cls(prover, farmer_public_key, pool_public_key, pool_contract_puzzle_hash, plot_public_key, time.time())

    def bump_last_use(self) -> None:
        self.last_use = time.time()

    def expired(self, expiry_seconds: int) -> bool:
        return time.time() - self.last_use > expiry_seconds


@dataclass
class Cache:
    _path: Path
    _changed: bool = False
    _data: dict[Path, CacheEntry] = field(default_factory=dict)
    expiry_seconds: int = 7 * 24 * 60 * 60  # Keep the cache entries alive for 7 days after its last access

    def __post_init__(self) -> None:
        self._path.parent.mkdir(parents=True, exist_ok=True)

    def __len__(self) -> int:
        return len(self._data)

    def update(self, path: Path, entry: CacheEntry) -> None:
        self._data[path] = entry
        self._changed = True

    def remove(self, cache_keys: list[Path]) -> None:
        for key in cache_keys:
            if key in self._data:
                del self._data[key]
                self._changed = True

    def save(self) -> None:
        try:
            disk_cache_entries: dict[str, DiskCacheEntry] = {
                str(path): DiskCacheEntry(
                    bytes(cache_entry.prover),
                    cache_entry.farmer_public_key,
                    cache_entry.pool_public_key,
                    cache_entry.pool_contract_puzzle_hash,
                    cache_entry.plot_public_key,
                    uint64(cache_entry.last_use),
                )
                for path, cache_entry in self.items()
            }
            cache_data: CacheDataV1 = CacheDataV1(
                [(plot_id, cache_entry) for plot_id, cache_entry in disk_cache_entries.items()]
            )
            disk_cache: VersionedBlob = VersionedBlob(uint16(CURRENT_VERSION), bytes(cache_data))
            serialized: bytes = bytes(disk_cache)
            self._path.write_bytes(serialized)
            self._changed = False
            log.info(f"Saved {len(serialized)} bytes of cached data")
        except Exception as e:
            log.error(f"Failed to save cache: {e}, {traceback.format_exc()}")

    def load(self) -> None:
        try:
            serialized = self._path.read_bytes()
            log.info(f"Loaded {len(serialized)} bytes of cached data")
            stored_cache: VersionedBlob = VersionedBlob.from_bytes(serialized)
            if stored_cache.version == CURRENT_VERSION:
                start = time.time()
                cache_data: CacheDataV1 = CacheDataV1.from_bytes(stored_cache.blob)
                self._data = {}
                estimated_c2_sizes: dict[int, int] = {}
                measured_sizes: dict[int, int] = {
                    32: 738,
                    33: 1083,
                    34: 1771,
                    35: 3147,
                    36: 5899,
                    37: 11395,
                    38: 22395,
                    39: 44367,
                }
                for path, cache_entry in cache_data.entries:
                    prover: ProverProtocol = get_prover_from_bytes(path, cache_entry.prover_data)
                    new_entry = CacheEntry(
                        prover,
                        cache_entry.farmer_public_key,
                        cache_entry.pool_public_key,
                        cache_entry.pool_contract_puzzle_hash,
                        cache_entry.plot_public_key,
                        float(cache_entry.last_use),
                    )
                    # TODO, drop the below entry dropping after few versions or whenever we force a cache recreation.
                    #       it's here to filter invalid cache entries coming from bladebit RAM plotting.
                    #       Related: - https://github.com/Chia-Network/chia-blockchain/issues/13084
                    #                - https://github.com/Chia-Network/chiapos/pull/337
                    param = new_entry.prover.get_param()
                    if param.size_v1 is not None:
                        k = param.size_v1
                        if k not in estimated_c2_sizes:
                            estimated_c2_sizes[k] = ceil(2**k / 100_000_000) * ceil(k / 8)
                        memo_size = len(new_entry.prover.get_memo())
                        prover_size = len(cache_entry.prover_data)
                        # Estimated C2 size + memo size + 2000 (static data + path)
                        # static data: version(2) + table pointers (<=96) + id(32) + k(1) => ~130
                        # path: up to ~1870, all above will lead to false positive.
                        # See https://github.com/Chia-Network/chiapos/blob/3ee062b86315823dd775453ad320b8be892c7df3/src/prover_disk.hpp#L282-L287  # noqa: E501

                        # Use experimental measurements if more than estimates
                        # https://github.com/Chia-Network/chia-blockchain/issues/16063
                        check_size = estimated_c2_sizes[k] + memo_size + 2000
                        if k in measured_sizes:
                            check_size = max(check_size, measured_sizes[k])

                        if prover_size > check_size:
                            log.warning(
                                "Suspicious cache entry dropped. Recommended: stop the harvester, remove "
                                f"{self._path}, restart. Entry: size {prover_size}, path {path}"
                            )
                        else:
                            self._data[Path(path)] = new_entry
                    elif param.strength_v2 is not None:
                        # TODO: todo_v2_plots validate prover size
                        self._data[Path(path)] = new_entry

                log.info(f"Parsed {len(self._data)} cache entries in {time.time() - start:.2f}s")

            else:
                raise ValueError(f"Invalid cache version {stored_cache.version}. Expected version {CURRENT_VERSION}.")
        except FileNotFoundError:
            log.debug(f"Cache {self._path} not found")
        except Exception as e:
            log.error(f"Failed to load cache: {e}, {traceback.format_exc()}")

    def keys(self) -> KeysView[Path]:
        return self._data.keys()

    def values(self) -> ValuesView[CacheEntry]:
        return self._data.values()

    def items(self) -> ItemsView[Path, CacheEntry]:
        return self._data.items()

    def get(self, path: Path) -> CacheEntry | None:
        return self._data.get(path)

    def changed(self) -> bool:
        return self._changed

    def path(self) -> Path:
        return self._path