emod-cli/debug_mod/DEBUG_ENV_SCRIPT/IPCSystem.py

# -*- coding: utf-8 -*-
import mod.server.extraServerApi as serverApi
import mod.client.extraClientApi as clientApi
from .Config import GET_DEBUG_IPC_PORT
import socket
import threading
import time
import json

def U16_BE(b):
    # type: (bytearray | str) -> int
    if isinstance(b, bytearray):
        return (b[0] << 8) | b[1]
    return (ord(b[0]) << 8) | ord(b[1])

def U32_BE(b):
    # type: (bytearray | str) -> int
    if isinstance(b, bytearray):
        return (b[0] << 24) | (b[1] << 16) | (b[2] << 8) | b[3]
    return (ord(b[0]) << 24) | (ord(b[1]) << 16) | (ord(b[2]) << 8) | ord(b[3])

class IPCSystem:
    def __init__(self, port=None):
        # type: (int | None) -> None
        self.port = port
        self.sock = None
        self.mLock = threading.Lock()
        self.handlers = {}
        self.pendingLock = threading.Lock()
        self.pendingClient = []
        self.pendingServer = []
        self.running = False
        self.generation = 0
        self.thread = None
        self.nextStartTime = 0.0

    def registerHandler(self, typeID, handler):
        # type: (int, callable) -> None
        self.handlers[typeID] = handler

    def updateHandlers(self, handlers):
        # type: (dict[int, callable]) -> None
        self.handlers.update(handlers)

    def start(self):
        now = time.time()
        with self.mLock:
            if self.running or not self.port or now < self.nextStartTime:
                return
            self.running = True
            self.nextStartTime = now + 1.0
            self.generation += 1
            generation = self.generation
            thread = threading.Thread(
                target=self._threadListenLoop,
                args=(generation,),
            )
            thread.daemon = True
            self.thread = thread
        thread.start()

    def close(self):
        with self.mLock:
            self.running = False
            self.generation += 1
            self.nextStartTime = 0.0
            sock = self.sock
            self.sock = None
            thread = self.thread
        if sock:
            try:
                sock.shutdown(socket.SHUT_RDWR)
            except socket.error:
                pass
            try:
                sock.close()
            except socket.error:
                pass
        if thread:
            thread.join(1.0)
        with self.pendingLock:
            self.pendingClient = []
            self.pendingServer = []

    def updateClient(self):
        self.start()
        with self.pendingLock:
            pending = self.pendingClient
            self.pendingClient = []
        self._runPending(pending)

    def updateServer(self):
        with self.pendingLock:
            pending = self.pendingServer
            self.pendingServer = []
        self._runPending(pending)

    @staticmethod
    def _runPending(pending):
        for handler, data in pending:
            try:
                handler(data)
            except Exception:
                import traceback
                traceback.print_exc()

    def _threadListenLoop(self, generation):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        with self.mLock:
            if not self.running or generation != self.generation:
                sock.close()
                return
            self.sock = sock

        try:
            sock.settimeout(2.0)
            sock.connect(("localhost", self.port))
            sock.settimeout(0.2)
            print("[IPCSystem] 已连接到调试服务器，端口：" + str(self.port))
            while self._isCurrent(generation, sock):
                header = self._recvAll(generation, sock, 6)
                typeID = U16_BE(header[0:2])
                dataLength = U32_BE(header[2:6])
                data = self._recvAll(generation, sock, dataLength)
                self._queueHandler(generation, typeID, data)
        except EOFError:
            pass
        except socket.error:
            if self._isCurrent(generation, sock):
                import traceback
                traceback.print_exc()
        except Exception:
            import traceback
            traceback.print_exc()
        finally:
            try:
                sock.close()
            except socket.error:
                pass
            with self.mLock:
                if self.sock is sock:
                    self.sock = None
                if self.thread is threading.currentThread():
                    self.thread = None
                if generation == self.generation:
                    self.running = False
            print("[IPCSystem] 连接已关闭")

    def _isCurrent(self, generation, sock):
        with self.mLock:
            return (
                self.running
                and generation == self.generation
                and self.sock is sock
            )

    def _queueHandler(self, generation, typeID, data):
        handler = self.handlers.get(typeID)
        if handler is None:
            print("[IPCSystem] 未知的TypeID数据包：" + str(typeID))
            return
        with self.mLock:
            if not self.running or generation != self.generation:
                return
            with self.pendingLock:
                if typeID == 4:
                    self.pendingServer.append((handler, data))
                else:
                    self.pendingClient.append((handler, data))

    def _recvAll(self, generation, sock, length):
        buf = bytearray()
        while len(buf) < length:
            try:
                more = sock.recv(length - len(buf))
            except socket.timeout:
                if not self._isCurrent(generation, sock):
                    raise EOFError("IPC connection stopped")
                continue
            if not more:
                raise EOFError("Socket closed before receiving all data")
            buf.extend(more)
        return buf

_CL_GAME_COMP = None
_SR_GAME_COMP = None

def AUTO_RELOAD(_=None):
    from .Game import RELOAD_MOD
    if _CL_GAME_COMP:
        _CL_GAME_COMP.AddTimer(0, lambda: RELOAD_MOD())
        return

def FAST_RELOAD(data):
    from .Game import RELOAD_ONCE_MODULE
    pathList = json.loads(str(data))
    def _FAST_RELOAD():
        for path in pathList:
            if RELOAD_ONCE_MODULE(path):
                print("[FAST_RELOAD] Reloaded module successfully: \"" + path + "\"")
    if _CL_GAME_COMP:
        _CL_GAME_COMP.AddTimer(0, _FAST_RELOAD)
        return

def EXEC_CLIENT_CODE(data):
    code = compile(str(data), "<string>", "exec")
    def _EXEC_CODE():
        print("[CLIENT_CODE] Executed successfully: " + str(eval(code)))
    _CL_GAME_COMP.AddTimer(0, _EXEC_CODE)

def EXEC_SERVER_CODE(data):
    code = compile(str(data), "<string>", "exec")
    def _EXEC_CODE():
        print("[SERVER_CODE] Executed successfully: " + str(eval(code)))
    _SR_GAME_COMP.AddTimer(0, _EXEC_CODE)

def RELOAD_GAME(_=None):
    def _RELOAD_GAME():
        from .Game import RELOAD_WORLD
        print("[RELOAD_GAME] Reloading the game...")
        RELOAD_WORLD()
    _CL_GAME_COMP.AddTimer(0, _RELOAD_GAME)

def RELOAD_SHADERS(_=None):
    def _RELOAD_SHADERS():
        from .Game import RELOAD_SHADERS
        RELOAD_SHADERS()
    _CL_GAME_COMP.AddTimer(0, _RELOAD_SHADERS)

def RELOAD_ONCE_SHADERS(fileName):
    def _RELOAD_ONCE_SHADERS():
        if clientApi.ReloadOneShader(str(fileName)):
            print("[RELOAD_ONCE_SHADERS] Reloaded shaders successfully.")
            return
        print("[RELOAD_ONCE_SHADERS] Failed to reload shaders.")
    _CL_GAME_COMP.AddTimer(0, _RELOAD_ONCE_SHADERS)

def RELOAD_ADDON_AND_GAME(_=None):
    def _RELOAD_ADDON_AND_GAME():
        from .Game import RELOAD_WORLD, RELOAD_ADDON
        print("[RELOAD_ADDON_AND_GAME] Reloading the addon and the game...")
        RELOAD_ADDON()
        RELOAD_WORLD()
    _CL_GAME_COMP.AddTimer(0, _RELOAD_ADDON_AND_GAME)

_IPCSYSTEM = IPCSystem(GET_DEBUG_IPC_PORT())
_IPCSYSTEM.updateHandlers(
    {
        1: AUTO_RELOAD,
        2: FAST_RELOAD,
        3: EXEC_CLIENT_CODE,
        4: EXEC_SERVER_CODE,
        5: RELOAD_GAME,
        6: RELOAD_SHADERS,
        7: RELOAD_ONCE_SHADERS,
        8: RELOAD_ADDON_AND_GAME,
    }
)

def UPDATE_CLIENT():
    _IPCSYSTEM.updateClient()

def UPDATE_SERVER():
    _IPCSYSTEM.updateServer()

def ON_CLIENT_INIT():
    global _CL_GAME_COMP
    _CL_GAME_COMP = clientApi.GetEngineCompFactory().CreateGame(clientApi.GetLevelId())

def ON_CLIENT_EXIT():
    global _CL_GAME_COMP
    _IPCSYSTEM.close()
    _CL_GAME_COMP = None

def ON_SERVER_INIT():
    global _SR_GAME_COMP
    _SR_GAME_COMP = serverApi.GetEngineCompFactory().CreateGame(serverApi.GetLevelId())

def ON_SERVER_EXIT():
    global _SR_GAME_COMP
    _SR_GAME_COMP = None