clan-core/pkgs/clan-app/clan_app/components/executor.py

import logging
import os
import signal
import sys
import traceback
from pathlib import Path
from typing import Any

import gi

gi.require_version("GdkPixbuf", "2.0")

import dataclasses
import multiprocessing as mp
from collections.abc import Callable

log = logging.getLogger(__name__)


# Kill the new process and all its children by sending a SIGTERM signal to the process group
def _kill_group(proc: mp.Process) -> None:
    pid = proc.pid
    if proc.is_alive() and pid:
        os.killpg(pid, signal.SIGTERM)
    else:
        log.warning(f"Process '{proc.name}' with pid '{pid}' is already dead")


@dataclasses.dataclass(frozen=True)
class MPProcess:
    name: str
    proc: mp.Process
    out_file: Path

    # Kill the new process and all its children by sending a SIGTERM signal to the process group
    def kill_group(self) -> None:
        _kill_group(proc=self.proc)


def _set_proc_name(name: str) -> None:
    if sys.platform != "linux":
        return
    import ctypes

    # Define the prctl function with the appropriate arguments and return type
    libc = ctypes.CDLL("libc.so.6")
    prctl = libc.prctl
    prctl.argtypes = [
        ctypes.c_int,
        ctypes.c_char_p,
        ctypes.c_ulong,
        ctypes.c_ulong,
        ctypes.c_ulong,
    ]
    prctl.restype = ctypes.c_int

    # Set the process name to "my_process"
    prctl(15, name.encode(), 0, 0, 0)


def _init_proc(
    func: Callable,
    out_file: Path,
    proc_name: str,
    on_except: Callable[[Exception, mp.process.BaseProcess], None] | None,
    **kwargs: Any,
) -> None:
    # Create a new process group
    os.setsid()

    # Open stdout and stderr
    with open(out_file, "w") as out_fd:
        os.dup2(out_fd.fileno(), sys.stdout.fileno())
        os.dup2(out_fd.fileno(), sys.stderr.fileno())

    # Print some information
    pid = os.getpid()
    gpid = os.getpgid(pid=pid)

    # Set the process name
    _set_proc_name(proc_name)

    # Close stdin
    sys.stdin.close()

    linebreak = "=" * 5
    # Execute the main function
    print(linebreak + f" {func.__name__}:{pid} " + linebreak, file=sys.stderr)
    try:
        func(**kwargs)
    except Exception as ex:
        traceback.print_exc()
        if on_except is not None:
            on_except(ex, mp.current_process())

            # Kill the new process and all its children by sending a SIGTERM signal to the process group
            pid = os.getpid()
            gpid = os.getpgid(pid=pid)
            print(f"Killing process group pid={pid} gpid={gpid}", file=sys.stderr)
            os.killpg(gpid, signal.SIGTERM)
        sys.exit(1)
    # Don't use a finally block here, because we want the exitcode to be set to
    # 0 if the function returns normally


def spawn(
    *,
    out_file: Path,
    on_except: Callable[[Exception, mp.process.BaseProcess], None] | None,
    func: Callable,
    **kwargs: Any,
) -> MPProcess:
    # Decouple the process from the parent
    if mp.get_start_method(allow_none=True) is None:
        mp.set_start_method(method="forkserver")

    # Set names
    proc_name = f"MPExec:{func.__name__}"

    # Start the process
    proc = mp.Process(
        target=_init_proc,
        args=(func, out_file, proc_name, on_except),
        name=proc_name,
        kwargs=kwargs,
    )
    proc.start()

    # Return the process
    mp_proc = MPProcess(name=proc_name, proc=proc, out_file=out_file)

    return mp_proc