From 03496bb28d8d057e29ced0dd549e95b491b8b21c Mon Sep 17 00:00:00 2001 From: R1kaB3rN <100738684+R1kaB3rN@users.noreply.github.com> Date: Mon, 18 Mar 2024 13:16:16 -0700 Subject: [PATCH] Add feature to set systemd as subreaper --- ULWGL/ulwgl_consts.py | 3 + ULWGL/ulwgl_plugins.py | 56 +++++++++++-- ULWGL/ulwgl_run.py | 93 +++++++++++++++++++-- ULWGL/ulwgl_test_plugins.py | 159 ++++++++++++++++++++++++++++++++++++ 4 files changed, 299 insertions(+), 12 deletions(-) diff --git a/ULWGL/ulwgl_consts.py b/ULWGL/ulwgl_consts.py index f3e58218..9f8ee776 100644 --- a/ULWGL/ulwgl_consts.py +++ b/ULWGL/ulwgl_consts.py @@ -1,5 +1,6 @@ from enum import Enum from pathlib import Path +from typing import Dict, Any class Color(Enum): @@ -35,3 +36,5 @@ class Color(Enum): "getcompatpath", "getnativepath", } + +TOMLDocument = Dict[str, Any] diff --git a/ULWGL/ulwgl_plugins.py b/ULWGL/ulwgl_plugins.py index 21900206..9cfae7bb 100644 --- a/ULWGL/ulwgl_plugins.py +++ b/ULWGL/ulwgl_plugins.py @@ -4,11 +4,13 @@ from typing import Dict, Set, Any, List, Tuple from argparse import Namespace from shutil import which +from ulwgl_log import log +from ulwgl_consts import TOMLDocument, ULWGL_LOCAL def set_env_toml( env: Dict[str, str], args: Namespace -) -> Tuple[Dict[str, str], List[str]]: +) -> Tuple[Dict[str, str], List[str], TOMLDocument]: """Read a TOML file then sets the environment variables for the Steam RT. In the TOML file, certain keys map to Steam RT environment variables. For example: @@ -54,11 +56,13 @@ def set_env_toml( opts = val elif key == "launch_args" and isinstance(val, str): opts = val.split(" ") + elif key == "reaper" and not val: + env["ULWGL_SYSTEMD"] = "1" - return env, opts + return env, opts, toml -def _check_env_toml(env: Dict[str, str], toml: Dict[str, Any]) -> Dict[str, Any]: +def _check_env_toml(env: Dict[str, str], toml: TOMLDocument) -> Dict[str, Any]: """Check for required or empty key/value pairs when reading a TOML config. NOTE: Casing matters in the config and we do not check if the game id is set @@ -92,12 +96,18 @@ def _check_env_toml(env: Dict[str, str], toml: Dict[str, Any]) -> Dict[str, Any] err: str = f"Value for key '{key}' in TOML is not a directory: {dir}" raise NotADirectoryError(err) - # Check for empty keys + # Check for empty and optional keys for key, val in toml[table].items(): + if key == "reaper" and not isinstance(val, bool): + err: str = ( + f"Value is not a boolean for '{key}' in TOML.\n" + f"Please specify a value or remove the entry:\n{key} = {val}" + ) + raise ValueError(err) if not val and isinstance(val, str): err: str = ( f"Value is empty for '{key}' in TOML.\n" - f"Please specify a value or remove the following entry:\n{key} = {val}" + f"Please specify a value or remove the entry:\n{key} = {val}" ) raise ValueError(err) @@ -189,3 +199,39 @@ def enable_zenity(command: str, opts: List[str], msg: str) -> None: # Close the Zenity process's standard input zenity_proc.stdin.close() zenity_proc.wait() + + +def enable_systemd(env: Dict[str, str], command: List[str]) -> List[str]: + """Use systemd to monitor and keep track of descendent processes. + + Descendent processes of ulwgl-run will be executed in a transient, user scoped unit + For more information of systemd-run, please visit + https://www.freedesktop.org/software/systemd/man/latest/systemd-run.html + """ + bin: str = which("systemd-run") + id: str = env["ULWGL_ID"] + + if not id.startswith("ulwgl-"): + id = "ulwgl-" + env["ULWGL_ID"] + + # Fallback to reaper + if not bin: + log.debug("systemd-run is not found in system\nUsing reaper as subreaper") + return enable_reaper( + env, + command, + ) + + # TODO Allow users to pass their own options + command.extend( + [ + bin, + "--user", + "--scope", + "--send-sighup", + "--description", + id, + ] + ) + + return command diff --git a/ULWGL/ulwgl_run.py b/ULWGL/ulwgl_run.py index de69097d..b965cb76 100755 --- a/ULWGL/ulwgl_run.py +++ b/ULWGL/ulwgl_run.py @@ -5,16 +5,28 @@ from argparse import ArgumentParser, Namespace, RawTextHelpFormatter from pathlib import Path from typing import Dict, Any, List, Set, Union, Tuple -from ulwgl_plugins import enable_steam_game_drive, set_env_toml, enable_reaper from re import match from subprocess import run from ulwgl_dl_util import get_ulwgl_proton -from ulwgl_consts import PROTON_VERBS, DEBUG_FORMAT, STEAM_COMPAT, ULWGL_LOCAL from ulwgl_util import setup_ulwgl from ulwgl_log import log, console_handler, CustomFormatter from ulwgl_util import UnixUser from logging import INFO, WARNING, DEBUG from errno import ENETUNREACH +from shutil import which +from json import loads as json_loads +from ulwgl_plugins import ( + enable_steam_game_drive, + set_env_toml, + enable_reaper, + enable_systemd, +from ulwgl_consts import ( + PROTON_VERBS, + DEBUG_FORMAT, + STEAM_COMPAT, + ULWGL_LOCAL, + TOMLDocument, +) def parse_args() -> Union[Namespace, Tuple[str, List[str]]]: # noqa: D103 @@ -207,11 +219,23 @@ def set_env( env["STEAM_COMPAT_TOOL_PATHS"] = env["PROTONPATH"] + ":" + ULWGL_LOCAL.as_posix() env["STEAM_COMPAT_MOUNTS"] = env["STEAM_COMPAT_TOOL_PATHS"] + # Gamescope + if "ULWGL_GAMESCOPE" in os.environ: + env["ULWGL_GAMESCOPE"] = os.environ["ULWGL_GAMESCOPE"] + + # Systemd + if "ULWGL_SYSTEMD" in os.environ: + env["ULWGL_SYSTEMD"] = os.environ["ULWGL_SYSTEMD"] + return env def build_command( - env: Dict[str, str], local: Path, command: List[str], opts: List[str] = None + env: Dict[str, str], + local: Path, + command: List[str], + opts: List[str] = None, + config: TOMLDocument = None, ) -> List[str]: """Build the command to be executed.""" verb: str = env["PROTON_VERB"] @@ -231,7 +255,20 @@ def build_command( err: str = "The following file was not found in PROTONPATH: proton" raise FileNotFoundError(err) - enable_reaper(env, command, local) + # Subreaper + if ( + config + and config.get("ulwgl").get("reaper") + and not config.get("ulwgl").get("reaper") + ): + log.debug("Using systemd as subreaper") + enable_systemd(env, command) + elif env.get("ULWGL_SYSTEMD") == "1": + log.debug("Using systemd as subreaper") + enable_systemd(env, command) + else: + log.debug("Using reaper as subreaper") + enable_reaper(env, command, local) command.extend([local.joinpath("ULWGL").as_posix(), "--verb", verb, "--"]) command.extend( @@ -270,6 +307,7 @@ def main() -> int: # noqa: D103 "STORE": "", "PROTON_VERB": "", "ULWGL_ID": "", + "ULWGL_SYSTEMD": "", } command: List[str] = [] opts: List[str] = None @@ -279,6 +317,7 @@ def main() -> int: # noqa: D103 # Expects this dir to be in sync with root # On update, files will be selectively updated args: Union[Namespace, Tuple[str, List[str]]] = parse_args() + config: TOMLDocument = None if "musl" in os.environ.get("LD_LIBRARY_PATH", ""): err: str = "This script is not designed to run on musl-based systems" @@ -316,7 +355,7 @@ def main() -> int: # noqa: D103 # Check environment if isinstance(args, Namespace) and getattr(args, "config", None): - env, opts = set_env_toml(env, args) + env, opts, config = set_env_toml(env, args) else: opts = args[1] # Reference the executable options check_env(env) @@ -337,17 +376,57 @@ def main() -> int: # noqa: D103 os.environ[key] = val # Run - build_command(env, ULWGL_LOCAL, command, opts) + build_command(env, ULWGL_LOCAL, command, opts, config) log.debug(command) return run(command).returncode +def stop_unit(id: str) -> None: + """Handle explicit kill or server shutdowns of the game. + + Intended to run when using gamescope with the launcher + + Used when systemd is configured as the subreaper and we assume the systemd + transient unit is ours if the unit's description matches the ULWGL ID + """ + if not os.environ["ULWGL_SYSTEMD"] == "1" and os.environ["ULWGL_GAMESCOPE"] == "1": + emoji: str = "\U0001f480" + log.warning("Explicit shutdown detected") + log.warning("Zombies will prevent re-running the game %s ...", emoji) + return + + result: str = run( + [which("systemctl"), "list-units", "--user", "-o", "json"], + capture_output=True, + text=True, + ).stdout + ulwgl_id: str = f"ulwgl-{id}" + unit: str = "" + + for item in json_loads(result): + if item.get("description") == ulwgl_id: + unit = item["unit"] + break + + if unit: + emoji: str = "\U0001f480" + log.console(f"Reaping zombies due to explicit shutdown {emoji} ...") + run([which("systemctl"), "stop", "--user", f"{unit}"]) + + if __name__ == "__main__": try: - sys.exit(main()) + ret: int = main() + + if not ret: + # Handle force exits when using gamescope + stop_unit(os.environ["ULWGL_ID"]) + + sys.exit(ret) except KeyboardInterrupt: log.warning("Keyboard Interrupt") + stop_unit(os.environ["ULWGL_ID"]) sys.exit(1) except SystemExit as e: if e.code != 0: diff --git a/ULWGL/ulwgl_test_plugins.py b/ULWGL/ulwgl_test_plugins.py index 6fb182ea..7d9a0d34 100644 --- a/ULWGL/ulwgl_test_plugins.py +++ b/ULWGL/ulwgl_test_plugins.py @@ -170,6 +170,127 @@ def tearDown(self): if self.test_local_share.exists(): rmtree(self.test_local_share.as_posix()) + def test_build_command_systemd(self): + """Test build_command when systemd is set as the subreaper.""" + test_toml = "foo.toml" + toml_str = f""" + [ulwgl] + prefix = "{self.test_file}" + proton = "{self.test_file}" + game_id = "{self.test_file}" + launch_args = ["{self.test_file}", "{self.test_file}"] + exe = "{self.test_exe}" + reaper = false + """ + toml_path = self.test_file + "/" + test_toml + result = None + test_command = [] + test_command_result = None + toml = None + + Path(self.test_file + "/proton").touch() + Path(toml_path).touch() + + with Path(toml_path).open(mode="w") as file: + file.write(toml_str) + + with patch.object( + ulwgl_run, + "parse_args", + return_value=argparse.Namespace(config=toml_path), + ): + # Args + result = ulwgl_run.parse_args() + # Config + self.env, opts, toml = ulwgl_plugins.set_env_toml(self.env, result) + # Prefix + ulwgl_run.setup_pfx(self.env["WINEPREFIX"]) + # Env + ulwgl_run.set_env(self.env, result) + # Game drive + ulwgl_plugins.enable_steam_game_drive(self.env) + + # Mock setting up the runtime + with patch.object( + ulwgl_util, + "setup_runtime", + return_value=None, + ): + ulwgl_util._install_ulwgl( + self.test_user_share, self.test_local_share, self.test_compat, json + ) + copytree( + Path(self.test_user_share, "sniper_platform_0.20240125.75305"), + Path(self.test_local_share, "sniper_platform_0.20240125.75305"), + dirs_exist_ok=True, + symlinks=True, + ) + copy(Path(self.test_user_share, "run"), Path(self.test_local_share, "run")) + copy( + Path(self.test_user_share, "run-in-sniper"), + Path(self.test_local_share, "run-in-sniper"), + ) + copy( + Path(self.test_user_share, "ULWGL"), + Path(self.test_local_share, "ULWGL"), + ) + # When the runtime updates, pressure vessel needs to be updated + copytree( + Path(self.test_user_share, "pressure-vessel"), + Path(self.test_local_share, "pressure-vessel"), + dirs_exist_ok=True, + symlinks=True, + ) + + for key, val in self.env.items(): + os.environ[key] = val + + # Build + test_command_result = ulwgl_run.build_command( + self.env, self.test_local_share, test_command, config=toml + ) + self.assertTrue( + test_command_result is test_command, "Expected the same reference" + ) + + # Verify contents of the command + ( + reaper, + systemd_opt0, + systemd_opt1, + systemd_opt2, + systemd_opt3, + systemd_opt4, + entry_point, + opt1, + verb, + opt2, + proton, + verb2, + exe, + ) = [*test_command] + # The entry point dest could change and binary paths could vary. + # Just check if there's a value for the reaper and the entry point + self.assertTrue(reaper, "Expected systemd") + self.assertEqual(systemd_opt0, "--user", "Expected --user") + self.assertEqual(systemd_opt1, "--scope", "Expected --scope") + self.assertEqual(systemd_opt2, "--send-sighup", "Expected --sighup") + self.assertEqual(systemd_opt3, "--description", "Expected --description") + self.assertEqual( + systemd_opt4, "ulwgl-" + self.env["ULWGL_ID"], "Expected ULWGL_ID to be set" + ) + self.assertTrue(entry_point, "Expected an entry point") + self.assertEqual(opt1, "--verb", "Expected --verb") + self.assertEqual(verb, self.test_verb, "Expected a verb") + self.assertEqual(opt2, "--", "Expected --") + self.assertEqual( + proton, + Path(self.env.get("PROTONPATH") + "/proton").as_posix(), + "Expected the proton file", + ) + self.assertEqual(verb2, self.test_verb, "Expected a verb") + self.assertEqual(exe, self.env["EXE"], "Expected the EXE") + def test_build_command_entry(self): """Test build_command. @@ -410,6 +531,44 @@ def test_build_command_toml(self): self.assertEqual(verb2, self.test_verb, "Expected a verb") self.assertEqual(exe, self.env["EXE"], "Expected the EXE") + def test_set_env_toml_reaper(self): + """Test set_env_toml for reaper. + + A ValueError should be raised if a boolean is not set + """ + test_toml = "foo.toml" + toml_path = self.test_file + "/" + test_toml + toml_str = f""" + [ulwgl] + prefix = "{self.test_file}" + proton = "{self.test_file}" + game_id = "{self.test_file}" + launch_args = ["{toml_path}"] + exe = "{self.test_exe}" + reaper = "foo" + """ + result = None + + Path(toml_path).touch() + + with Path(toml_path).open(mode="w") as file: + file.write(toml_str) + + with patch.object( + ulwgl_run, + "parse_args", + return_value=argparse.Namespace(config=toml_path), + ): + # Args + result = ulwgl_run.parse_args() + self.assertIsInstance( + result, Namespace, "Expected a Namespace from parse_arg" + ) + self.assertTrue(vars(result).get("config"), "Expected a value for --config") + # Env + with self.assertRaisesRegex(ValueError, "reaper"): + ulwgl_plugins.set_env_toml(self.env, result) + def test_set_env_toml_nofile(self): """Test set_env_toml for values that are not a file.