Added more folders and restructured classes inside them.
This commit is contained in:
287
src/ptdiag/experiments/cleanup.py
Normal file
287
src/ptdiag/experiments/cleanup.py
Normal file
@@ -0,0 +1,287 @@
|
||||
"""Experiment cleanup, prefix-link rollback, and registry restoration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from ..config import AppConfig
|
||||
from ..prefix import cleanup_dedicated_wineserver
|
||||
from ..registry import DesktopRegistrySnapshot
|
||||
from ..result import ResultDirectory
|
||||
from ..runner import CommandLogger, CommandResult
|
||||
from .models import ExperimentError
|
||||
|
||||
|
||||
EXPLORER_KEY = r"HKCU\Software\Wine\Explorer"
|
||||
DESKTOPS_KEY = r"HKCU\Software\Wine\Explorer\Desktops"
|
||||
DISPOSABLE_DIRECTORY_NAME = "ptdiag-disposable"
|
||||
|
||||
ProtonCommand = Callable[..., list[str]]
|
||||
WineserverCleaner = Callable[..., CommandResult]
|
||||
|
||||
|
||||
def replace_game_link(
|
||||
prefix_link: Path,
|
||||
disposable_game: Path,
|
||||
expected_live_game: Path,
|
||||
logger: CommandLogger,
|
||||
) -> str:
|
||||
"""Point the dedicated prefix at a verified disposable game tree."""
|
||||
|
||||
if not prefix_link.is_symlink():
|
||||
raise ExperimentError(
|
||||
f"Dedicated-prefix game path is not a symlink: {prefix_link}"
|
||||
)
|
||||
if prefix_link.resolve(strict=True) != expected_live_game.resolve(strict=True):
|
||||
raise ExperimentError("Dedicated-prefix game link changed after source preflight")
|
||||
original_target = os.readlink(prefix_link)
|
||||
require_command(
|
||||
logger.run(
|
||||
["ln", "-sfn", "--", str(disposable_game), str(prefix_link)],
|
||||
cwd=prefix_link.parent,
|
||||
timeout_seconds=15.0,
|
||||
),
|
||||
"point dedicated prefix at disposable game",
|
||||
)
|
||||
if (
|
||||
not prefix_link.is_symlink()
|
||||
or prefix_link.resolve(strict=True) != disposable_game.resolve(strict=True)
|
||||
):
|
||||
raise ExperimentError(
|
||||
"Dedicated-prefix game link did not resolve to disposable copy"
|
||||
)
|
||||
return original_target
|
||||
|
||||
|
||||
def restore_game_link(
|
||||
prefix_link: Path,
|
||||
original_target: str,
|
||||
expected_live_game: Path,
|
||||
logger: CommandLogger,
|
||||
) -> None:
|
||||
"""Restore the dedicated-prefix game link and verify its exact target."""
|
||||
|
||||
require_command(
|
||||
logger.run(
|
||||
["ln", "-sfn", "--", original_target, str(prefix_link)],
|
||||
cwd=prefix_link.parent,
|
||||
timeout_seconds=15.0,
|
||||
),
|
||||
"restore dedicated-prefix game link",
|
||||
)
|
||||
if not prefix_link.is_symlink():
|
||||
raise ExperimentError("Restored dedicated-prefix game path is not a symlink")
|
||||
if os.readlink(prefix_link) != original_target:
|
||||
raise ExperimentError("Dedicated-prefix game link text was not restored exactly")
|
||||
if prefix_link.resolve(strict=True) != expected_live_game.resolve(strict=True):
|
||||
raise ExperimentError("Dedicated-prefix game link target was not restored")
|
||||
|
||||
|
||||
def remove_disposable_game(
|
||||
config: AppConfig,
|
||||
result: ResultDirectory,
|
||||
disposable_game: Path,
|
||||
logger: CommandLogger,
|
||||
*,
|
||||
directory_name: str = DISPOSABLE_DIRECTORY_NAME,
|
||||
) -> None:
|
||||
"""Remove only the result-scoped disposable tree beneath its managed root."""
|
||||
|
||||
staging_root = (config.paths.app_root / directory_name).absolute()
|
||||
validate_disposable_path(
|
||||
staging_root,
|
||||
disposable_game,
|
||||
result.path.name,
|
||||
directory_name=directory_name,
|
||||
)
|
||||
require_command(
|
||||
logger.run(
|
||||
["rm", "-rf", "--one-file-system", "--", str(disposable_game)],
|
||||
timeout_seconds=900.0,
|
||||
),
|
||||
"remove disposable game tree",
|
||||
)
|
||||
if disposable_game.exists() or disposable_game.is_symlink():
|
||||
raise ExperimentError(
|
||||
f"Disposable game cleanup left its target: {disposable_game}"
|
||||
)
|
||||
try:
|
||||
staging_root.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def validate_disposable_path(
|
||||
staging_root: Path,
|
||||
path: Path,
|
||||
expected_name: str,
|
||||
*,
|
||||
directory_name: str = DISPOSABLE_DIRECTORY_NAME,
|
||||
) -> None:
|
||||
"""Reject cleanup targets outside the managed, result-scoped staging root."""
|
||||
|
||||
if path.parent != staging_root or path.name != expected_name:
|
||||
raise ExperimentError(f"Unsafe disposable game path: {path}")
|
||||
if staging_root.name != directory_name:
|
||||
raise ExperimentError(f"Unsafe disposable staging root: {staging_root}")
|
||||
|
||||
|
||||
def require_command(outcome: CommandResult, operation: str) -> None:
|
||||
if not command_ok(outcome):
|
||||
raise ExperimentError(
|
||||
f"Could not {operation} ({outcome.returncode}): "
|
||||
f"{outcome.error or outcome.stderr.strip() or 'command failed'}"
|
||||
)
|
||||
|
||||
|
||||
def registry_setup_commands(
|
||||
config: AppConfig,
|
||||
desktop_name: str,
|
||||
desktop_size: str,
|
||||
*,
|
||||
proton_command: ProtonCommand,
|
||||
) -> list[list[str]]:
|
||||
return [
|
||||
proton_command(
|
||||
config,
|
||||
"runinprefix",
|
||||
"reg.exe",
|
||||
"add",
|
||||
EXPLORER_KEY,
|
||||
"/v",
|
||||
"Desktop",
|
||||
"/t",
|
||||
"REG_SZ",
|
||||
"/d",
|
||||
desktop_name,
|
||||
"/f",
|
||||
),
|
||||
proton_command(
|
||||
config,
|
||||
"runinprefix",
|
||||
"reg.exe",
|
||||
"add",
|
||||
DESKTOPS_KEY,
|
||||
"/v",
|
||||
desktop_name,
|
||||
"/t",
|
||||
"REG_SZ",
|
||||
"/d",
|
||||
desktop_size,
|
||||
"/f",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def registry_restore_commands(
|
||||
config: AppConfig,
|
||||
snapshot: DesktopRegistrySnapshot,
|
||||
*,
|
||||
proton_command: ProtonCommand,
|
||||
) -> list[list[str]]:
|
||||
return [
|
||||
registry_restore_command(
|
||||
config,
|
||||
EXPLORER_KEY,
|
||||
"Desktop",
|
||||
snapshot.selected_desktop.exists,
|
||||
snapshot.selected_desktop.value,
|
||||
proton_command=proton_command,
|
||||
),
|
||||
registry_restore_command(
|
||||
config,
|
||||
DESKTOPS_KEY,
|
||||
snapshot.desktop_name,
|
||||
snapshot.desktop_size.exists,
|
||||
snapshot.desktop_size.value,
|
||||
proton_command=proton_command,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def registry_restore_command(
|
||||
config: AppConfig,
|
||||
key: str,
|
||||
name: str,
|
||||
existed: bool,
|
||||
value: str | None,
|
||||
*,
|
||||
proton_command: ProtonCommand,
|
||||
) -> list[str]:
|
||||
if existed:
|
||||
assert value is not None
|
||||
return proton_command(
|
||||
config,
|
||||
"runinprefix",
|
||||
"reg.exe",
|
||||
"add",
|
||||
key,
|
||||
"/v",
|
||||
name,
|
||||
"/t",
|
||||
"REG_SZ",
|
||||
"/d",
|
||||
value,
|
||||
"/f",
|
||||
)
|
||||
return proton_command(
|
||||
config,
|
||||
"runinprefix",
|
||||
"reg.exe",
|
||||
"delete",
|
||||
key,
|
||||
"/v",
|
||||
name,
|
||||
"/f",
|
||||
)
|
||||
|
||||
|
||||
def run_checked_commands(
|
||||
logger: CommandLogger,
|
||||
commands: Sequence[Sequence[str]],
|
||||
*,
|
||||
cwd: Path,
|
||||
env_overrides: Mapping[str, str],
|
||||
) -> None:
|
||||
for command in commands:
|
||||
outcome = logger.run(
|
||||
command,
|
||||
cwd=cwd,
|
||||
env_overrides=env_overrides,
|
||||
timeout_seconds=15.0,
|
||||
)
|
||||
if not command_ok(outcome):
|
||||
raise ExperimentError(
|
||||
f"Command failed ({outcome.returncode}): {list(command)!r}: "
|
||||
f"{outcome.error or outcome.stderr.strip()}"
|
||||
)
|
||||
|
||||
|
||||
def cleanup_checked(
|
||||
config: AppConfig,
|
||||
logger: CommandLogger,
|
||||
*,
|
||||
cleaner: WineserverCleaner = cleanup_dedicated_wineserver,
|
||||
) -> None:
|
||||
"""Stop the dedicated wineserver while accepting its idle return contract."""
|
||||
|
||||
outcome = cleaner(config, logger, timeout_seconds=15.0)
|
||||
if (
|
||||
outcome.returncode == 1
|
||||
and not outcome.timed_out
|
||||
and outcome.error is None
|
||||
and not outcome.stdout
|
||||
and not outcome.stderr
|
||||
):
|
||||
return
|
||||
if not command_ok(outcome):
|
||||
raise ExperimentError(
|
||||
"Dedicated-prefix wineserver cleanup failed: "
|
||||
f"{outcome.error or outcome.stderr.strip() or outcome.returncode}"
|
||||
)
|
||||
|
||||
|
||||
def command_ok(result: CommandResult) -> bool:
|
||||
return not result.timed_out and result.error is None and result.returncode == 0
|
||||
Reference in New Issue
Block a user