chia-blockchain/chia/util/files.py

import asyncio
import logging
import os
import shutil

from aiofiles import tempfile  # type: ignore
from pathlib import Path
from typing import Union


log = logging.getLogger(__name__)


def move_file(src: Path, dst: Path):
    """
    Attempts to move the file at src to dst, falling back to a copy if the move fails.
    """

    dir_perms: int = 0o700
    # Create the parent directory if necessary
    os.makedirs(dst.parent, mode=dir_perms, exist_ok=True)

    try:
        # Attempt an atomic move first
        os.replace(os.fspath(src), os.fspath(dst))
    except Exception as e:
        log.debug(f"Failed to move {src} to {dst} using os.replace, reattempting with shutil.move: {e}")
        try:
            # If that fails, use the more robust shutil.move(), though it may internally initiate a copy
            shutil.move(os.fspath(src), os.fspath(dst))
        except Exception:
            log.exception(f"Failed to move {src} to {dst} using shutil.move")
            raise


async def move_file_async(src: Path, dst: Path, *, reattempts: int = 6, reattempt_delay: float = 0.5):
    """
    Attempts to move the file at src to dst, making multiple attempts if the move fails.
    """

    remaining_attempts: int = reattempts
    while True:
        try:
            move_file(src, dst)
        except Exception:
            if remaining_attempts > 0:
                log.debug(f"Failed to move {src} to {dst}, retrying in {reattempt_delay} seconds")
                remaining_attempts -= 1
                await asyncio.sleep(reattempt_delay)
            else:
                break
        else:
            break

    if not dst.exists():
        raise FileNotFoundError(f"Failed to move {src} to {dst}")
    else:
        log.debug(f"Moved {src} to {dst}")


async def write_file_async(file_path: Path, data: Union[str, bytes], *, file_mode: int = 0o600, dir_mode: int = 0o700):
    """
    Writes the provided data to a temporary file and then moves it to the final destination.
    """

    # Create the parent directory if necessary
    os.makedirs(file_path.parent, mode=dir_mode, exist_ok=True)

    mode: str = "w+" if type(data) == str else "w+b"
    temp_file_path: Path
    async with tempfile.NamedTemporaryFile(dir=file_path.parent, mode=mode, delete=False) as f:
        temp_file_path = f.name
        await f.write(data)
        await f.flush()
        os.fsync(f.fileno())

    try:
        await move_file_async(temp_file_path, file_path)
    except Exception:
        log.exception(f"Failed to move temp file {temp_file_path} to {file_path}")
    else:
        os.chmod(file_path, file_mode)
    finally:
        # We expect the file replace/move to have succeeded, but cleanup the temp file just in case
        try:
            if Path(temp_file_path).exists():
                os.remove(temp_file_path)
        except Exception:
            log.exception(f"Failed to remove temp file {temp_file_path}")