ykxmz/windows/lanremote/instance.py

# -*- coding: utf-8 -*-
"""单实例控制。

窗口关掉后程序仍在后台跑，用户再次双击 EXE 时应该能**找回窗口**，
而不是又起一个进程（那样端口会冲突，用户会以为坏了）。

做法：主实例占住一个本机回环端口；后启动的实例连上去发个 "show" 就退出。

端口可以用环境变量 YKXMZ_LOCK_PORT 覆盖：
    测试需要在"用户自己的实例正在跑"的同时再起一个被控端。
    不加这个口子的话，测试实例一启动就看到锁被占，直接退出
    （日志里只有一句"已有实例在运行，退出"），
    表现成"测试怎么都不登记到中转"，很难查。
    默认值不变，正常运行完全不受影响。
"""
from __future__ import annotations

import os
import socket
import threading

CONTROL_PORT = 8767
_CMD_SHOW = b"SHOW"


def control_port() -> int:
    """单实例锁用的端口（允许用环境变量覆盖，主要给测试用）。"""
    raw = str(os.environ.get("YKXMZ_LOCK_PORT") or "").strip()
    if raw:
        try:
            value = int(raw)
            if 1 <= value <= 65535:
                return value
        except ValueError:
            pass
    return CONTROL_PORT


class InstanceLock:
    def __init__(self, port: int | None = None):
        self.port = control_port() if port is None else int(port)
        self._sock: socket.socket | None = None
        self._thread: threading.Thread | None = None
        self._running = False
        self.on_show = None            # 收到"显示窗口"请求时的回调

    @property
    def is_primary(self) -> bool:
        return self._sock is not None

    def try_claim(self) -> bool:
        """尝试成为主实例。失败说明已有实例在运行。"""
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            # 关键：Windows 上 SO_REUSEADDR 会允许第二个进程绑同一个端口，
            # 导致单实例判定失效。Windows 必须用 SO_EXCLUSIVEADDRUSE。
            if hasattr(socket, "SO_EXCLUSIVEADDRUSE"):
                try:
                    sock.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
                except OSError:
                    pass
            else:
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            sock.bind(("127.0.0.1", self.port))
            sock.listen(4)
            sock.settimeout(1.0)
        except OSError:
            try:
                sock.close()
            except Exception:
                pass
            return False

        self._sock = sock
        self._running = True
        self._thread = threading.Thread(target=self._loop, name="instance",
                                        daemon=True)
        self._thread.start()
        return True

    def _loop(self) -> None:
        while self._running:
            try:
                conn, _addr = self._sock.accept()
            except socket.timeout:
                continue
            except OSError:
                break
            try:
                conn.settimeout(2.0)
                data = conn.recv(32)
                if data.strip() == _CMD_SHOW and self.on_show:
                    try:
                        self.on_show()
                    except Exception:
                        pass
                conn.sendall(b"OK")
            except OSError:
                pass
            finally:
                try:
                    conn.close()
                except OSError:
                    pass

    def release(self) -> None:
        self._running = False
        if self._thread is not None and self._thread.is_alive():
            self._thread.join(timeout=2.0)
        self._thread = None
        if self._sock is not None:
            try:
                self._sock.close()
            except OSError:
                pass
            self._sock = None


def is_running(port: int | None = None, timeout: float = 0.6) -> bool:
    """现在有没有别的实例占着单实例锁。

    只是"探一下"，不会让对方做任何事（发的是空包，不是 SHOW）。

    为什么单独要这个函数：
        "双击 exe 开设置窗口"这条路上要知道后台是不是在跑 ——
        在跑的话，保存完得让后台重启才会生效；
        不在跑的话直接存就行。原来没有这个能力，
        我第一版顺手写了个 InstanceLock.probe_free()，
        而那个方法根本不存在，一调就 AttributeError。
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(timeout)
    try:
        sock.connect(("127.0.0.1", control_port() if port is None else int(port)))
        return True
    except OSError:
        return False
    finally:
        try:
            sock.close()
        except OSError:
            pass


def signal_show(port: int | None = None, timeout: float = 2.0) -> bool:
    """通知已在运行的实例把窗口显示出来。"""
    if port is None:
        port = control_port()
    try:
        sock = socket.create_connection(("127.0.0.1", port), timeout=timeout)
        sock.sendall(_CMD_SHOW)
        try:
            sock.recv(16)
        except OSError:
            pass
        sock.close()
        return True
    except OSError:
        return False