Newer
Older
safe-algo-pro / apis / control.py
import os
import platform
import subprocess
import sys
import threading
import time
import traceback

from fastapi import APIRouter
import docker
import socket

from apis.base import standard_error_response, standard_response
from common.global_logger import logger

router = APIRouter()


def is_running_in_docker():
    """ 检测当前程序是否在 Docker 容器中运行 """
    try:
        with open('/proc/1/cgroup', 'rt') as f:
            return 'docker' in f.read()
    except Exception:
        return False


def is_windows_host():
    try:
        # 尝试解析 host.docker.internal
        socket.gethostbyname("host.docker.internal")
        return True
    except socket.error:
        return False


def get_container_id():
    try:
        with open('/proc/self/cgroup', 'r') as f:
            for line in f:
                if 'docker' in line:
                    # 解析出容器 ID
                    container_id = line.strip().split('/')[-1]
                    return container_id
    except Exception as e:
        logger.error(f"无法获取容器ID: {e}")
    return None


@router.get("/restart")
async def restart():
    try:
        # 立即返回响应的函数
        def restart_container_async():
            try:
                if is_running_in_docker():
                    if is_windows_host():
                        client = docker.DockerClient(base_url='tcp://host.docker.internal:2375')
                        container_id = client.containers.get(socket.gethostname())
                    else:
                        client = docker.DockerClient(base_url='unix://var/run/docker.sock')
                        container_id = get_container_id()

                    # 重启容器
                    container = client.containers.get(container_id)
                    container.restart()
                    logger.info("Container restarted successfully.")
                else:
                    os_type = platform.system()

                    # 获取当前脚本路径和命令行参数
                    command = [sys.executable] + sys.argv

                    if os_type == "Windows":
                        # Windows 环境下的重启逻辑
                        subprocess.Popen(["start", "python"] + sys.argv, shell=True)
                        time.sleep(1)  # 确保新进程启动
                        sys.exit(0)  # 退出当前进程

                    elif os_type == "Linux" or os_type == "Darwin":
                        # Linux 和 macOS 环境下的重启逻辑
                        full_command = f"nohup {' '.join(command)} &"
                        subprocess.Popen(full_command, shell=True, stdout=subprocess.DEVNULL,
                                         stderr=subprocess.DEVNULL)
                        time.sleep(1)  # 确保新进程启动
                        sys.exit(0)  # 退出当前进程

                    else:
                        print(f"Unsupported OS for restart: {os_type}")
            except Exception as ex:
                traceback.print_exc()
                logger.error(f"Failed to restart container asynchronously: {ex}")

        # 在新线程中启动重启操作
        threading.Thread(target=restart_container_async).start()

        return standard_response()
    except Exception as e:
        traceback.print_exc()
        return standard_error_response(code=500, message=f"Failed to restart container: {e}")

@router.get("/sync_test")
def sync_test():
    return standard_response()

@router.get("/async_test")
async def async_test():
    return standard_response()