python-dmon 0.2.0__py3-none-any.whl

dmon/control.py

@@ -0,0 +1,542 @@
+import os
+from pathlib import Path
+import shlex
+import shutil
+import sys
+import subprocess
+import time
+from typing import List
+
+import psutil
+from termcolor import colored
+
+from .constants import DEFAULT_META_DIR, META_SUFFIX
+from .types import DmonTaskConfig, DmonMeta, PathType
+from .utils import len_ansi, pad_ansi
+
+
+def ensure_meta_dir(meta_path: Path):
+    """
+    Ensure the meta directory exists.
+    If it is under the default meta directory, create .gitignore to ignore all files in it.
+    """
+    meta_dir = meta_path.resolve().parent
+    meta_dir.mkdir(parents=True, exist_ok=True)
+    if meta_dir == DEFAULT_META_DIR.resolve():
+        gitignore = DEFAULT_META_DIR / ".gitignore"
+        if not gitignore.exists():
+            gitignore.write_text("# Auto-generated by dmon\n*\n")
+
+
+def ensure_log_dir(log_path: Path):
+    log_path.parent.mkdir(parents=True, exist_ok=True)
+
+
+def start(cfg: DmonTaskConfig):
+    meta_path = Path(cfg.meta_path).resolve()
+    log_path = Path(cfg.log_path).resolve()
+    cwd = Path(cfg.cwd).resolve()
+
+    try:
+        ret_meta = DmonMeta.load(meta_path)
+    except Exception:
+        ret_meta = None
+    if ret_meta:
+        print(
+            f"{colored('Start failed: meta file already exists', color='red', attrs=['bold'])}",
+            file=sys.stderr,
+        )
+        print_status(ret_meta)
+        print(
+            "\nRun 'dmon status' / 'dmon list' to check, or 'dmon stop' to stop it.",
+            file=sys.stderr,
+        )
+        return 1
+
+    ensure_meta_dir(meta_path)
+    ensure_log_dir(log_path)
+
+    env = None  # default behavior of Popen
+    if cfg.override_env:
+        env = cfg.env
+    elif cfg.env:
+        env = {**os.environ, **cfg.env}
+
+    shell = isinstance(cfg.cmd, str)
+
+    # Platform-specific parameters to run the process in background detached from parent
+    kwargs = {}
+    if sys.platform.startswith("win"):
+        # DETACHED_PROCESS = 0x00000008
+        CREATE_NO_WINDOW = 0x08000000
+        kwargs["creationflags"] = CREATE_NO_WINDOW
+        if isinstance(cfg.cmd, list) and len(cfg.cmd) > 0:
+            # On Windows, use full path for the executable when shell=False
+            cfg.cmd[0] = shutil.which(cfg.cmd[0]) or cfg.cmd[0]
+    else:
+        # Make the child process independent of the parent process in Unix-like systems
+        kwargs["start_new_session"] = True
+
+    meta = DmonMeta(
+        name=cfg.name,
+        meta_path=str(meta_path),
+        log_path=str(log_path),
+        log_rotate=cfg.log_rotate,
+        cmd=cfg.cmd,
+        cwd=str(cwd),
+        env=cfg.env,
+        override_env=cfg.override_env,
+        shell=shell,
+        popen_kwargs=kwargs,
+    )
+
+    if cfg.log_rotate:
+        rotate_log_path = Path(cfg.rotate_log_path).resolve()
+
+        meta.rotate_log_path = str(rotate_log_path)
+        meta.log_max_size = cfg.log_max_size
+        meta.rotate_log_max_size = cfg.rotate_log_max_size
+
+        ensure_log_dir(rotate_log_path)
+
+        # use runner to start user process and handle log rotation
+        if isinstance(cfg.cmd, str):
+            # when cmd is str, we need to split it into list for subprocess
+            cmd = shlex.split(cfg.cmd)
+        else:
+            cmd = cfg.cmd
+        args = [
+            sys.executable,
+            "-m",
+            "dmon.runner",
+            "--log-path",
+            str(log_path),
+            "--max-log-size",
+            str(cfg.log_max_size),
+            "--rotate-log-path",
+            str(rotate_log_path),
+            "--max-rotate-log-size",
+            str(cfg.rotate_log_max_size),
+        ]
+        if shell:
+            args.append("--shell")
+        args.append("--")
+        args.extend(cmd)
+        proc = subprocess.Popen(
+            args,
+            cwd=cwd,
+            env=env,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.STDOUT,
+            **kwargs,
+        )
+    else:
+        # Open the log file (append binary mode)
+        with open(log_path, "ab", buffering=0) as lof:
+            # Start the child process with stdout/stderr redirected to the log
+            proc = subprocess.Popen(
+                cfg.cmd,
+                stdout=lof,
+                stderr=subprocess.STDOUT,
+                cwd=cwd,
+                env=env,
+                shell=shell,
+                text=False,  # binary mode
+                bufsize=0,  # unbuffered
+                **kwargs,
+            )
+
+    meta.pid = proc.pid
+    try:
+        p = psutil.Process(proc.pid)
+        create_time = p.create_time()
+        create_time_human = time.strftime(
+            "%Y-%m-%d %H:%M:%S", time.localtime(create_time)
+        )
+        meta.create_time = create_time
+        meta.create_time_human = create_time_human
+    except psutil.NoSuchProcess:
+        # process already exited?
+        pass
+
+    meta.dump(meta_path)
+
+    print_status(meta)
+    return 0
+
+
+def stop(
+    meta_path: PathType,
+    timeout=5.0,
+):
+    meta_path = Path(meta_path).resolve()
+    meta = DmonMeta.load(meta_path)
+    if meta is None:
+        print(
+            colored(
+                "Stop failed: meta file not found (maybe not started)\n",
+                color="red",
+                attrs=["bold"],
+            )
+            + str(meta_path),
+            file=sys.stderr,
+        )
+        return 1
+
+    pid = meta.pid
+    try:
+        proc = psutil.Process(pid)
+    except psutil.NoSuchProcess:
+        print(
+            colored(
+                f"Process {pid} not found (already exited); removing stale meta file",
+                color="yellow",
+                attrs=["bold"],
+            ),
+            file=sys.stderr,
+        )
+        print_status(meta)
+        meta_path.unlink(missing_ok=True)
+        return 1
+
+    # check if it's the same process by comparing create_time
+    if not check_same_process(proc, meta.create_time):
+        print(
+            colored(
+                f"Warning: PID {pid} exists but create_time does not match (maybe reused by another process); removing stale meta file",
+                color="yellow",
+                attrs=["bold"],
+            ),
+            file=sys.stderr,
+        )
+        print_status(meta)
+        meta_path.unlink(missing_ok=True)
+        return 1
+
+    if sys.platform.startswith("win"):
+        ret = terminate_win(proc, timeout)
+    else:
+        ret = terminate_posix(proc, timeout)
+    print_status(meta)
+    if ret == 0:
+        meta_path.unlink(missing_ok=True)
+    return ret
+
+
+def terminate_posix(proc: psutil.Process, timeout):
+    # send SIGTERM first for graceful shutdown (if platform supports)
+    proc.terminate()
+
+    # wait for the process to exit or timeout
+    try:
+        ret = proc.wait(timeout)
+        print(
+            colored(
+                f"Process {proc.pid} exited with code {ret}; removing meta file",
+                color="green",
+                attrs=["bold"],
+            ),
+            file=sys.stderr,
+        )
+    except psutil.TimeoutExpired:
+        print(
+            colored(
+                f"Process {proc.pid} did not exit in time; shutting down child processes",
+                color="yellow",
+                attrs=["bold"],
+            ),
+            file=sys.stderr,
+        )
+
+        # first shut down child processes (leaf nodes first)
+        for child in reversed(proc.children(recursive=True)):
+            try:
+                child.terminate()
+                child.wait(timeout=2)
+            except Exception:
+                try:
+                    child.kill()
+                except Exception:
+                    pass
+        # then shut down the parent process
+        try:
+            # check if process is already exited
+            ret = proc.wait(timeout=2)
+            print(
+                colored(
+                    f"Process {proc.pid} exited with code {ret} after terminating children; removing meta file",
+                    color="green",
+                    attrs=["bold"],
+                ),
+                file=sys.stderr,
+            )
+        except psutil.TimeoutExpired:
+            print(
+                colored(
+                    f"Process {proc.pid} did not exit in time after terminating children; killing it",
+                    color="yellow",
+                    attrs=["bold"],
+                ),
+                file=sys.stderr,
+            )
+            try:
+                proc.kill()
+                print(
+                    colored(
+                        f"Killed process {proc.pid}; removing meta file",
+                        color="green",
+                        attrs=["bold"],
+                    ),
+                    file=sys.stderr,
+                )
+            except Exception as e:
+                print(
+                    colored(
+                        f"Failed to kill process {proc.pid}: {e}",
+                        color="red",
+                        attrs=["bold"],
+                    ),
+                    file=sys.stderr,
+                )
+                return 1
+    return 0
+
+
+def terminate_win(proc: psutil.Process, timeout):
+    # On Windows, we need to stop child processes first
+    children = proc.children(recursive=True)
+    # reverse to kill leaf nodes first
+    for child in reversed(children):
+        try:
+            child.kill()
+            child.wait(timeout=2)
+        except Exception:
+            pass
+    try:
+        ret = proc.wait(timeout=timeout)
+        print(
+            colored(
+                f"Process {proc.pid} exited with code {ret}; removing meta file",
+                color="green",
+                attrs=["bold"],
+            ),
+            file=sys.stderr,
+        )
+    except psutil.TimeoutExpired:
+        print(
+            colored(
+                f"Process {proc.pid} did not exit in time after killing children; killing it",
+                color="yellow",
+                attrs=["bold"],
+            ),
+            file=sys.stderr,
+        )
+        proc.kill()
+        print(
+            colored(
+                f"Killed process {proc.pid}; removing meta file",
+                color="green",
+                attrs=["bold"],
+            ),
+            file=sys.stderr,
+        )
+    return 0
+
+
+def restart(
+    cfg: DmonTaskConfig,
+    timeout=5.0,
+):
+    stop(cfg.meta_path, timeout)
+    print("--- Restarting ---", file=sys.stderr)
+    return start(cfg)
+
+
+def status(meta_path: PathType):
+    meta_path = Path(meta_path).resolve()
+    meta = DmonMeta.load(meta_path)
+    if meta is None:
+        print(
+            colored(
+                "Status failed: meta file not found (maybe not started)\n",
+                color="red",
+                attrs=["bold"],
+            )
+            + str(meta_path),
+            file=sys.stderr,
+        )
+        return 1
+
+    print_status(meta)
+    print("\nProcess Tree:", file=sys.stderr)
+    print_process_table([meta])
+    return 0
+
+
+def check_same_process(proc: psutil.Process, create_time: float) -> bool:
+    if create_time < 0:
+        return False
+    return abs(proc.create_time() - create_time) < 1e-3
+
+
+def get_unique_process(pid: int, create_time: float):
+    """
+    Get the process with given PID and create_time if it is still the same process.
+    """
+    if create_time >= 0:
+        try:
+            p = psutil.Process(pid)
+            if check_same_process(p, create_time):
+                return p
+        except psutil.NoSuchProcess:
+            pass
+    return None
+
+
+def check_running(pid: int, create_time: float) -> bool:
+    """
+    Check if a process with given PID and create_time is running.
+    """
+    # return psutil.pid_exists(pid)
+    if create_time < 0:
+        return False
+    try:
+        p = psutil.Process(pid)
+        return abs(p.create_time() - create_time) < 1e-3
+    except psutil.NoSuchProcess:
+        return False
+
+
+def print_status(meta: DmonMeta):
+    status = (
+        colored("Running", on_color="on_green")
+        if check_running(meta.pid, meta.create_time)
+        else colored("Exited", on_color="on_light_red")
+    )
+
+    # key-value pairs with aligned keys
+    rows = [
+        ("TASK", colored(meta.name, "cyan", attrs=["bold"])),
+        ("PID", colored(str(meta.pid), "cyan", attrs=["bold"])),
+        ("STATUS", status),
+        ("CMD", meta.cmd),
+        ("WORKING DIR", meta.cwd),
+        ("CREATE TIME", meta.create_time_human),
+        ("META PATH", meta.meta_path),
+        ("LOG ROTATE", meta.log_rotate),
+        ("LOG PATH", meta.log_path),
+    ]
+    if meta.log_rotate:
+        rows.append(("ROTATE LOG PATH", meta.rotate_log_path))
+        rows.append(("LOG MAX SIZE", f"{meta.log_max_size} MB"))
+        rows.append(("ROTATE LOG MAX SIZE", f"{meta.rotate_log_max_size} MB"))
+
+    # calculate the max width of the keys
+    key_width = max(len(key) for key, _ in rows)
+    # align the keys and print
+    lines = [
+        f"{colored(f'{key:<{key_width}}: ', attrs=['dark'])}{value}"
+        for key, value in rows
+    ]
+    print("\n".join(lines), file=sys.stderr)
+
+
+def get_table_row(p: psutil.Process, target_ppid: int, prefix=""):
+    ppid = p.ppid()
+    ppid_str = (
+        colored(str(ppid), "cyan", attrs=["bold"]) if ppid == target_ppid else str(ppid)
+    )
+    return (
+        prefix + p.name(),
+        p.pid,
+        ppid_str,
+        p.status(),
+        p.cmdline(),
+        time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(p.create_time())),
+        "N/A",
+    )
+
+
+def print_process_table(metas: List[DmonMeta], full_width: bool = False):
+    headers = ("TASK", "PID", "PPID", "STATUS", "CMD", "CREATE TIME", "LOG PATH")
+    align = ("<", ">", ">", "<", "<", "<", "<")
+
+    # add table header
+    rows = []
+    rows.append(headers)
+
+    processes = []
+    for meta in metas:
+        proc = get_unique_process(meta.pid, meta.create_time)
+        if proc:
+            status = colored("Running", on_color="on_green")
+            ppid = proc.ppid()
+            processes.append(proc)
+        else:
+            status = colored("Exited", on_color="on_light_red")
+            ppid = "N/A"
+        rows.append(
+            (
+                colored(meta.name, "cyan", attrs=["bold"]),
+                colored(meta.pid, "cyan", attrs=["bold"]),
+                ppid,
+                status,
+                meta.cmd,
+                meta.create_time_human,
+                meta.log_path,
+            )
+        )
+        # add child processes indented
+        if proc:
+            children = proc.children(recursive=True)
+            for idx, child in enumerate(children):
+                prefix = "├ " if idx < len(children) - 1 else "└ "
+                rows.append(get_table_row(child, target_ppid=proc.pid, prefix=prefix))
+                processes.append(child)
+
+    # calculate column widths
+    widths = [max(len_ansi(str(row[i])) for row in rows) for i in range(len(headers))]
+
+    term_width = shutil.get_terminal_size().columns
+    diff = sum(widths) + 2 * (len(headers) - 1) - term_width
+    if diff > 0 and not full_width:
+        # truncate CMD column
+        cmd_idx = headers.index("CMD")
+        widths[cmd_idx] = max(widths[cmd_idx] - diff, 10)  # at least 10 chars
+
+    # print the table with proper padding
+    lines = []
+    for row in rows:
+        line = "  ".join(
+            # pad the cell except the last one if left-aligned (no need for extra spaces)
+            (
+                pad_ansi(str(cell), widths[i], align[i])
+                if i < len(align) - 1 or align[i] != "<"
+                else str(cell)
+            )
+            for i, cell in enumerate(row)
+        )
+        lines.append(line)
+    print("\n".join(lines), file=sys.stderr)
+
+    return processes
+
+
+def list_processes(dir: PathType, full_width: bool):
+    target_dmon_dir = Path(dir).resolve()
+    metas = []
+    if target_dmon_dir.exists() and target_dmon_dir.is_dir():
+        for meta_file in target_dmon_dir.glob(f"*{META_SUFFIX}"):
+            meta = DmonMeta.load(meta_file)
+            if meta is not None:
+                metas.append(meta)
+    # sort by name (case-insensitive)
+    metas.sort(key=lambda m: m.name.lower())
+    n_task = len(metas)
+    processes = print_process_table(metas, full_width)
+    n_proc = len(processes)
+    print(
+        f"\nFound {n_task} task{'s' if n_task > 1 else ''} ({n_proc} process{'es' if n_proc > 1 else ''}) in {target_dmon_dir}",
+        file=sys.stderr,
+    )
+    return 0