mirror of
https://github.com/Chia-Network/chia-blockchain.git
synced 2026-08-29 18:14:26 -05:00
Matt HauffandGitHub
40db4635a8
* Enable PEP604 Ruff rules * Fix harcoded signature in test * Hack CLVMStreamable test with note to fast follow
97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
# Package: utils
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
|
|
from aiofiles import tempfile
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def move_file(src: Path, dst: Path) -> None:
|
|
"""
|
|
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) -> None:
|
|
"""
|
|
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: str | bytes, *, file_mode: int = 0o600, dir_mode: int = 0o700
|
|
) -> None:
|
|
"""
|
|
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: Literal["w+", "w+b"] = "w+" if type(data) is str else "w+b"
|
|
temp_file_path: Path
|
|
async with tempfile.NamedTemporaryFile(dir=file_path.parent, mode=mode, delete=False) as f:
|
|
# Ignoring type error since it is not obvious how to tie the type of the data
|
|
# being passed in to the type of the file object, etc.
|
|
temp_file_path = f.name # type: ignore[assignment]
|
|
await f.write(data) # type: ignore[arg-type]
|
|
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}")
|