scripts/main.py

import os
import sys
import logging
import asyncio
from typing import List
from concurrent.futures import ThreadPoolExecutor
from alibabacloud_ecs20140526.client import Client as Ecs20140526Client
from alibabacloud_alb20200616.client import Client as Alb20200616Client
from alibabacloud_tea_openapi import models as open_api_models
from alibabacloud_ecs20140526 import models as ecs_20140526_models
from alibabacloud_alb20200616 import models as alb_20200616_models
from alibabacloud_tea_util import models as util_models
from alibabacloud_tea_util.client import Client as UtilClient
import config
import utils

class Sample:
    def __init__(self):
        pass

    @staticmethod
    def create_ecs_client() -> Ecs20140526Client:
        config = open_api_models.Config(  # 创建Config对象，用于配置访问凭证
            access_key_id=config.ALIBABA_CLOUD_ACCESS_KEY_ID,  # 从config中获取ALIBABA_CLOUD_ACCESS_KEY_ID参数
            access_key_secret=config.ALIBABA_CLOUD_ACCESS_KEY_SECRET  # 从config中获取ALIBABA_CLOUD_ACCESS_KEY_SECRET参数
        )
        config.endpoint = 'ecs-cn-hangzhou.aliyuncs.com'  # 设置ECS服务的访问地址
        return Ecs20140526Client(config)  # 返回配置好的ECS客户端对象

    @staticmethod
    def create_alb_client() -> Alb20200616Client:
        config = open_api_models.Config(  # 创建Config对象，用于配置访问凭证
            access_key_id=config.ALIBABA_CLOUD_ACCESS_KEY_ID,  # 从config中获取ALIBABA_CLOUD_ACCESS_KEY_ID参数
            access_key_secret=config.ALIBABA_CLOUD_ACCESS_KEY_SECRET  # 从config中获取ALIBABA_CLOUD_ACCESS_KEY_SECRET参数
        )
        config.endpoint = 'alb.cn-hangzhou.aliyuncs.com'  # 设置ALB服务的访问地址
        return Alb20200616Client(config)  # 返回配置好的ALB客户端对象

    @staticmethod
    def create_ecs_instance(ecs_client: Ecs20140526Client, amount: int) -> List[str]:
        instance_ids = []
        for i in range(amount):
            request = ecs_20140526_models.CreateInstanceRequest(  # 创建ECS实例请求
                size=config.instance_config['size'],  # 实例大小
                category=config.instance_config['category'],  # 实例类型
                region_id=config.instance_config['region_id'],  # 地域ID
                image_id=config.instance_config['image_id'],  # 镜像ID
                instance_type=config.instance_config['instance_type'],  # 实例类型
                security_group_id=config.instance_config['security_group_id'],  # 安全组ID
                v_switch_id=config.instance_config['v_switch_id'],  # 虚拟交换机ID
                instance_name=f"{config.instance_config['instance_name']}-{i+1}",  # 实例名称
                unique_suffix=config.instance_config['unique_suffix'],  # 是否使用唯一后缀
                password_inherit=config.instance_config['password_inherit'],  # 是否继承密码
                zone_id=config.instance_config['zone_id'],  # 可用区ID
                key_pair_name=config.instance_config['key_pair_name']  # 密钥对名称
            )
            response = ecs_client.create_instance(request)  # 调用ECS客户端的创建实例方法
            instance_id = response.body.instance_id  # 获取创建的实例ID
            instance_ids.append(instance_id)  # 将实例ID添加到列表中
            logging.info(f"Created ECS instance with ID: {instance_id}")  # 记录日志
        return instance_ids  # 返回实例ID列表

    @staticmethod
    async def start_services_and_health_check(create_client: Ecs20140526Client, slb_client: Alb20200616Client, ess_instance_ids: List[str]):
        try:
            await utils.send_file_to_ecs(client=create_client, instance_id_list=ess_instance_ids, **config.start_sh)  # 发送启动脚本到机器上
            logging.info(f"send start shell file finished, instances: {ess_instance_ids}")

            start_sh_param = "latest"
            server_start_sh = os.path.join(config.start_sh['target_dir'], config.start_sh['name'])
            server_start_commend = f"sh {server_start_sh} {start_sh_param}"
            await utils.run_command(client=create_client, instance_ids=ess_instance_ids, command=server_start_commend)  # 启动服务

            global health_instances
            health_instances = []
            max_wait_time = 180
            loop = asyncio.get_running_loop()
            max_workers = 10  # 设置最大线程数
            executor = ThreadPoolExecutor(max_workers=max_workers)
            tasks = [
                loop.run_in_executor(executor, utils.health_check, *args) for args in
                [(slb_client, instance_id, max_wait_time) for instance_id in ess_instance_ids]
            ]
            await asyncio.wait(tasks)  # 异步探活
            logging.info(f"health instances count: {len(health_instances)}, {health_instances}")

            if len(health_instances) > 0:  # 如果有机器探活成功
                time.sleep(20)

            Sample.mount_alb_to_instance(slb_client, ess_instance_ids)  # 将机器挂载到ALB服务器组中
        except Exception as e:
            logging.error(f"An error occurred: {e}")
            sys.exit(1)  # 退出程序并返回错误状态码

    @staticmethod
    def mount_alb_to_instance(client: Alb20200616Client, instance_ids: List[str]) -> None:
        try:
            servers = [alb_20200616_models.AddServersToServerGroupRequestServers(  # 创建要添加的服务器列表
                server_id=instance_id,
                server_type=config.alb_config['server_type'],
                weight=config.alb_config['weight']
            ) for instance_id in instance_ids]
            add_servers_to_server_group_request = alb_20200616_models.AddServersToServerGroupRequest(  # 创建添加服务器到服务器组的请求
                server_group_id=config.alb_config['server_group_id'],  # 服务器组ID
                servers=servers  # 服务器列表
            )
            runtime = util_models.RuntimeOptions()  # 创建运行时选项
            client.add_servers_to_server_group_with_options(add_servers_to_server_group_request, runtime)  # 将机器挂载到ALB服务器组中
        except Exception as e:
            logging.error(f"An error occurred: {e}")
            sys.exit(1)  # 退出程序并返回错误状态码

    @staticmethod
    def main(args: List[str]) -> None:
        try:
            ecs_client = Sample.create_ecs_client()  # 创建ECS客户端
            alb_client = Sample.create_alb_client()  # 创建ALB客户端
            amount = int(args[0]) if len(args) > 0 else 1  # 从命令行参数获取amount值，如果没有传入则默认为1
            instance_ids = Sample.create_ecs_instance(ecs_client, amount)  # 创建ECS实例并获取实例ID列表

            asyncio.run(Sample.start_services_and_health_check(ecs_client, alb_client, instance_ids))  # 使用asyncio运行启动服务和异步探活
        except Exception as e:
            logging.error(f"An error occurred: {e}")
            sys.exit(1)  # 退出程序并返回错误状态码

if __name__ == '__main__':
    Sample.main(sys.argv[1:])  # 调用main方法，传递命令行参数