vcast_license_usage_server/license.py

#!/usr/bin/python3
import json
import re
import socket
import asyncio
import subprocess
import threading
import time
from typing import Dict, List, Optional, Union
from datetime import datetime
from my_log import get_logger
import socket
import time
import threading
import select
log = get_logger()


MOCK_LMUTIL = True
if MOCK_LMUTIL:
    from mock_lmutil import *

DISPATCHER_SERVER_PORT = 8809

class LicenseDispatcher:
    def __init__(self, host='0.0.0.0', port=DISPATCHER_SERVER_PORT):
        self.host = host
        self.port = port
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(5)
        self.inputs = [self.server_socket]  # all socket to listen
        self.outputs = []  # all socket to write
        self.message_queues = {}  # store message to send
        self.lock = threading.Lock()
        self.wait_queue = Queue()
        self.exit_event = threading.Event()
        self.loop_thread = threading.Thread(target=self.loop)
        self.loop_thread.start()

    def exit(self):
        self.exit_event.set()
        self.loop_thread.join()

    def loop(self):
        log.info(f"dispatcher server start, listen {self.host}:{self.port}")
        while not self.exit_event.is_set():
            # select listen all socket read/write event
            readable, writable, exceptional = select.select(self.inputs, self.outputs, self.inputs, 1)
            # handle readable socket
            for sock in readable:
                if sock is self.server_socket:
                    # new client connect
                    self._handle_new_connection()
                else:
                    # handle client data
                    self._handle_client_data(sock)
            # handle writable socket
            for sock in writable:
                self._handle_client_response(sock)
            # handle exception socket
            for sock in exceptional:
                self._handle_socket_exception(sock)
            # clean inactive client
            self._cleanup_inactive_clients()
        log.debug("LicenseDispatcher exit")

    def _handle_new_connection(self):
        with self.lock:
            client_socket, address = self.server_socket.accept()
            log.info(f"new client connect: {address}")
            client_socket.setblocking(0)
            self.inputs.append(client_socket)
            self.message_queues[client_socket] = []

    def _handle_client_data(self, sock):
        try:
            data = sock.recv(1024)
            if not data:
                return
            log.trace(f"recv data: {data}")
            try:
                req = json.loads(data.decode('utf-8'))
            except json.JSONDecodeError:
                log.error(f"Invalid JSON data received: {data}")
                self._close_connection(sock)
                return
            req_type = req.get('type', None)
            if req_type == 'heartbeat':
                self._handle_heartbeat(sock, req)
            elif req_type == 'lic_received_ack':
                self._handle_lic_received_ack(req)
        except Exception as e:
            log.error(f"handle client data error: {e}")
            self._close_connection(sock)

    def _handle_heartbeat(self, sock, req):
        """ 心跳包内容
            {
                "type":"heartbeat",
                "user":"",
                "host":"",
                "heartbeat_time":datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }
        """
        for item in self.wait_queue:
            if item['user'] == req['user']:
                item['sock'] = sock
                item['heartbeat_time'] = datetime.now()
                break
        else:
            log.warning(f"user {req['user']} not in wait queue")
            sock.close()
            return
        # 如果用户已经分配到license，则直接返回
        if item['lic'] != '':
            self.send_lic_dispatch(item)
            return
        response = {
            'type':'heartbeat_ack',
            'status': 'ok',
            'heartbeat_time': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
        with self.lock:
            self.message_queues[sock].append(json.dumps(response))
            if sock not in self.outputs:
                self.outputs.append(sock)

    def _handle_lic_received_ack(self, req):
        self._remove_user(req['user'])

    def _handle_client_response(self, sock):
        try:
            with self.lock:
                if sock in self.message_queues and self.message_queues[sock]:
                    next_msg = self.message_queues[sock].pop(0)
                    sock.send(next_msg.encode('utf-8'))
                else:
                    self.outputs.remove(sock)
        except Exception as e:
            log.error(f"send response error: {e}")
            self._close_connection(sock)

    def _handle_socket_exception(self, sock):
        log.error(f"handle exception socket: {sock}")
        self._close_connection(sock)

    def _remove_user(self, user):
        with self.lock:
            for idx, item in enumerate(self.wait_queue):
                if item['user'] != user:
                    continue
                self.wait_queue.remove(idx)
                self._close_connection(item['sock'])
                break

    def _close_connection(self, sock):
        if not sock:
            return
        with self.lock:
            # 先从所有列表中移除socket
            if sock in self.outputs:
                self.outputs.remove(sock)
            if sock in self.inputs:
                self.inputs.remove(sock)
            if sock in self.message_queues:
                del self.message_queues[sock]
        try:
            sock.shutdown(socket.SHUT_RDWR)  # 先关闭连接
        except:
            pass  # 忽略已关闭socket的错误
        finally:
            sock.close()  # 最后关闭socket

    def _cleanup_inactive_clients(self, timeout=3):
        """clean inactive client"""
        current_time = datetime.now()
        remove_list = []
        with self.lock:
            for idx, item in enumerate(self.wait_queue):
                if (current_time - item['heartbeat_time']).total_seconds() < timeout:
                    continue
                log.info(f"clean inactive client: {item['user']}")
                remove_list.append(item['user'])
        for user in remove_list:
            self._remove_user(user)

    def add_wait_queue(self, user, host):
        data = {
            "user":user,
            "host":host,
            'status':'wait_dispatch',
            "start_time":datetime.now(),
            'heartbeat_time':datetime.now(),
            'lic':'',
            'sock':None
        }
        with self.lock:
            self.wait_queue.put(data)

    def dispatch_license(self, lic, available_cnt):
        """
        分配license存在几种情况：
        1. 该用户在等待队列中且没有分配到license，那么就直接分配
        2. 该用户在等待队列中且已经分配到license，那么就跳过
        3. 该lic已经被用户占用了，但是这个用户还没有连接，lic处于已经预分配状态，所以也不能再派发给别人
        """
        resp_list = []
        with self.lock:
            for item in self.wait_queue:
                if available_cnt <= 0:
                    break
                if item['lic'] == lic: # 如果该lic已经被分配了，可用数量减1
                    available_cnt -= 1
                    continue
                if item['status'] != 'wait_dispatch': # 如果该用户无需进行分配，则跳过
                    continue
                if item['lic'] != '': # 如果该用户已经分配了，则本次分配跳过
                    continue
                item['lic'] = lic
                item['status'] = 'dispatched'
                resp_list.append(item)
                available_cnt -= 1
        for item in resp_list:
            self.send_lic_dispatch(item)

    def send_lic_dispatch(self, item):
        sock = item['sock'] #如果sock不存在说明该用户还没连接，等待其连接后再发送lic
        if not sock:
            return
        response = {
            'user':item['user'],
            'host':item['host'],
            'lic':item['lic'],
            'type':'lic_dispatch',
            'status': 'ok',
            'msg':f'{item["user"]} get {item["lic"]} license'
        }
        with self.lock:
            self.message_queues[sock].append(json.dumps(response))
            if sock not in self.outputs:
                self.outputs.append(sock)


def _get_ip_address(domain):
    if MOCK_LMUTIL:
        return mock_ipaddress(domain)
    try:
        ip_address = socket.gethostbyname(domain)
        return ip_address
    except socket.gaierror as e:
        log.error(f"get ip address {domain} error: {e}")
        return None

def _get_valid_licenses():
    lic_id = 0
    while True:
        lic_id += 1
        domain = f"szmaslic{lic_id:02d}"
        if _get_ip_address(domain) is None:
            break
        yield "27003@" + domain

class UsedInfo:
    def __init__(self):
        self.user = ''
        self.host = ''
        self.lic = ''
        self.start_time:datetime = None
    def __str__(self):
        obj = self.as_dict()
        return str(obj)
    def __repr__(self):
        return self.__str__()
    def as_dict(self):
        time = datetime.now() - self.start_time
        login_time = (time.days * 24 + time.seconds / 3600.0)
        return {
            "user":self.user,
            "host":self.host,
            "lic":self.lic,
            "login":f'{login_time:.2f}h'
        }

class LicenseInfo:
    def __init__(self, name: str = '', total: int = 0, used: int = 0):
        self.name = name
        self.total = total
        self.used = used
        self.used_info:List[UsedInfo] = []

    def available(self):
        return self.total - self.used

    def __str__(self):
        obj = self.as_dict()
        return str(obj)
    def __repr__(self):
        return self.__str__()

    def as_dict(self):
        info = []
        for used in self.used_info:
            info.append(used.as_dict())
        return {
            "name":self.name,
            "total":self.total,
            "used":self.used,
            "used_info":info
        }

class Queue:
    def __init__(self):
        self.queue = []  # 使用列表作为底层存储
        self.lock = threading.Lock()  # 创建一个锁

    def put(self, item):
        with self.lock:
            self.queue.append(item)  # 添加元素到队列

    def pop(self):
        with self.lock:
            if not self.is_empty():
                return self.queue.pop(0)  # 从队列头部移除并返回元素
            else:
                return None  # 如果队列为空，返回None
    def get(self, idx):
        with self.lock:
            if not self.is_empty():
                return self.queue[idx]
            else:
                return None
    def __getitem__(self, idx):
        with self.lock:
            return self.queue[idx]
    def is_empty(self):
        with self.lock:
            return len(self.queue) == 0  # 检查队列是否为空

    def qsize(self):
        with self.lock:
            return len(self.queue)  # 返回队列大小

    def remove(self, idx):
        with self.lock:
            if 0 <= idx < len(self.queue):
                return self.queue.pop(idx)
            return None

class LicenseManager:
    def __init__(self):
        self.licenses_stat = LicenseStatistic()
        self.licenses:Dict[str,LicenseInfo] = {}
        self.dispatcher = LicenseDispatcher()
        self.exit_event = threading.Event()
        self.timer = None
        self.start_updater()  # 替换原来的线程启动

    def start_updater(self):
        if not self.exit_event.is_set():
            self.update_once()
            # 设置下一次执行
            self.timer = threading.Timer(5.0, self.start_updater)
            self.timer.start()

    def update_once(self):
        """单次更新操作"""
        self.licenses_stat.update()
        if self.dispatcher.wait_queue.qsize() > 0:
            for lic, available_cnt, _ in self.licenses_stat.available():
                self.dispatcher.dispatch_license(lic, available_cnt)

    def request(self, user: str, host: str) -> Dict[str, Union[str, dict]]:
        #如果已经是在线的状态就不允许申请
        rep = {
                "user":user,
                "host":host,
                "status":"refuse",
                "lic":'',
                "msg":f"unknown error"
        }
        if self.dispatcher.wait_queue.qsize() > 100:
                rep["status"] = "refuse"
                rep["msg"] = f"{user} queue is full!!"
                return rep
        if user in self.licenses_stat.online_users():
            rep["msg"] = f"{user} is online"
            log.warning(rep["msg"])
            return rep

        #如果是已经在队列中则返回等待
        item = next((item for item in self.dispatcher.wait_queue if item['user'] == user), None)
        if item:
            rep["status"] = "queue"
            rep["msg"] = f"{user} is in queue!!"
            rep["start_time"] = item['start_time'].strftime("%Y-%m-%d %H:%M:%S")
            return rep
        log.info(f"add {user} to wait queue")
        # 如果队列为空，则直接分配一个license
        if self.dispatcher.wait_queue.qsize() == 0:
            self.dispatcher.add_wait_queue(user, host)
            lic, available_cnt, _ = next(self.licenses_stat.available(), ('', 0, 0))
            self.dispatcher.dispatch_license(lic, available_cnt)
        else:
            self.dispatcher.add_wait_queue(user, host)
        rep["status"] = "wait_dispatch"
        rep["msg"] = f"{user} add to queue"
        rep["dispatcher_server"] = {
            "port": DISPATCHER_SERVER_PORT
        }
        return rep

    def get_wait_queue(self):
        resp = []
        for item in self.dispatcher.wait_queue:
            resp.append({
                "user":item['user'],
                "host":item['host'],
                "status":item['status'],
                "wait_time":int((datetime.now() - item['start_time']).total_seconds())
            })
        return resp

    def get_licenses_info(self):
        resp = {}
        for lic, info in self.licenses_stat.licenses.items():
            info  = info.as_dict()
            resp[lic] = info
        return resp

    def exit(self):
        self.exit_event.set()
        if self.timer:
            self.timer.cancel()
        self.timer = None
        if self.dispatcher:
            self.dispatcher.exit()

class LicenseStatistic:
    def __init__(self):
        self.licenses:Dict[str,LicenseInfo] = {}

    def update(self):
        # 更新license信息，使用异步获取
        tasks = []
        for lic in _get_valid_licenses():
            log.trace(f'find valid lic: {lic}')
            tasks.append(get_license_used_info(lic))
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            infos = loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
            for info in infos:
                if info is None:
                    continue
                if not isinstance(info, LicenseInfo):
                    log.error(f"get_license_used_info {info} is not LicenseInfo")
                    continue
                self.licenses[info.name] = info
            if not self.licenses:
                log.warning("not vcast lic found")
        except Exception as e:
            log.error(f"update license info error: {e}")
        loop.close()

    def get_lic_list(self):
        # 获取所有的license名称
        return self.licenses.keys()

    def available(self):
        # 获取所有的可用数量
        for _,lic in self.licenses.items():
            value = lic.available()
            if value > 0:
                yield (lic.name, value, lic.total) #lic的name和个数

    def online_users(self):
        # 获取在线用户
        for used_info in self.used():
            yield used_info.user #返回的是user name

    def used(self):
        # 获取所有的已使用信息
        for _,info in self.licenses.items():
            for used in info.used_info:
                yield used #返回的是UsedInfo对象

    def sorted(self):
        info = {}
        for used_info in self.used():
            time = datetime.now() - used_info.start_time
            login_time = (time.days * 24 + time.seconds / 3600.0)
            info[used_info.user] = login_time
        sorted_info = dict(sorted(info.items(), key=lambda item: item[1], reverse=True))
        return sorted_info #返回的是字典，key为username value为登录时间

    def __str__(self):
        string = str(self.licenses).replace("'", '"')
        log.info(string)
        obj = json.loads(string)
        string = json.dumps(obj, indent = 4)
        return string
    def __repr__(self):
        return self.__str__()


def shell(cmd, checkret=True):
    if MOCK_LMUTIL:
        ret = mock_lmutil(cmd)
        if not ret:
            return 1, '', ''
        return 0, ret, ''
    process = subprocess.Popen(cmd,
                             shell=True,
                             stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE,
                             universal_newlines=True)
    out, err = process.communicate()
    ret = process.returncode
    if checkret and ret != 0:
        raise subprocess.CalledProcessError(ret, cmd, output=out, stderr=err)
    return ret, out, err

"""
  {"27003@szmaslic03": {"name": "27003@szmaslic03", "total": 2, "used": 1, "used_info": "[{"name": "haoxiang.ran", "host": "szl3bc12808", "lic": "szmaslic03", "login": "0.41"}]"}, "27003@szmaslic04": {"name": "27003@szmaslic04", "total": 2, "used": 1, "used_info": "[{"name": "macro.yang", "host": "szl3bc12804", "lic": "szmaslic04", "login": "52.28"}]"}, "27003@szmaslic06": {"name": "27003@szmaslic06", "total": 2, "used": 2, "used_info": "[{"name": "zhenhao.he", "host": "szl3bc12804", "lic": "szmaslic06", "login": "0.20"}, {"name": "phillip.chan", "host": "szl3bc12810", "lic": "szmaslic06", "login": "0.16"}]"}, "27003@szmaslic07": {"name": "27003@szmaslic07", "total": 2, "used": 1, "used_info": "[{"name": "zhiyi.li", "host": "szl3bc12806", "lic": "szmaslic07", "login": "2.96"}]"}, "27003@szmaslic08": {"name": "27003@szmaslic08", "total": 2, "used": 2, "used_info": "[{"name": "soo.liu", "host": "szl3bc12809", "lic": "szmaslic08", "login": "6.55"}, {"name": "haichao.ou", "host": "szl3bc12806", "lic": "szmaslic08", "login": "5.93"}]"}, "27003@szmaslic09": {"name": "27003@szmaslic09", "total": 2, "used": 2, "used_info": "[{"name": "louie.liang", "host": "szl3bc12810", "lic": "szmaslic09", "login": "8.31"}, {"name": "zabbix", "host": "szl3bc06409", "lic": "szmaslic09", "login": "0.90"}]"}, "27003@szmaslic10": {"name": "27003@szmaslic10", "total": 2, "used": 2, "used_info": "[{"name": "harvey.li", "host": "szl3bc12810", "lic": "szmaslic10", "login": "57.38"}, {"name": "kw.hu", "host": "szl3bc12809", "lic": "szmaslic10", "login": "1.41"}]"}}
"""
async def get_license_used_info(lic):
    # log.info(f"get_license_used_info {lic}")
    server_used_info = f"timeout 1 /tools/software/vcast/flexlm/lmutil lmstat -a -c {lic}"
    ret,out,err = shell(server_used_info, checkret=False)
    # log.info(f"get_license_used_info {lic} {out} {ret} {err}")
    if ret != 0:
        return None
    """ data like:
    Users of VCAST_C_ENT_0800:  (Total of 2 licenses issued;  Total of 2 licenses in use)

    "VCAST_C_ENT_0800" v23, vendor: vector, expiry: permanent(no expiration date)
    vendor_string: CUST:64925:
    floating license

        zhenhao.he szl3bc12804 /dev/tty (v23) (szmaslic06/27003 337), start Thu 1/2 8:52
        phillip.chan szl3bc12810 /dev/tty (v23) (szmaslic06/27003 233), start Thu 1/2 14:59

    Users of CCAST_0801:  (Total of 2 licenses issued;  Total of 0 licenses in use)

    """
    status = 'idle'
    info = LicenseInfo(lic)
    for line in out.splitlines():
        line = line.strip()
        if status == 'idle':
            #Error getting status: Cannot find license file. (-1,359:2 "No such file or directory")
            if line.startswith("Error getting status"):
                log.warning(f"{lic}")
                return None
            if not line.startswith("Users of VCAST_C_ENT_0800:"):
                continue
            status = 'start'
            #[Users of VCAST_C_ENT_0800:  (Total of 2 licenses issued;  Total of 1 license in use)
            pattern = r"\(Total of (\d+).*Total of (\d+).*\)"
            ret = re.findall(pattern, line)
            if not ret:
                log.warning(f"{lic} parser [{line}] fail")
                return info
            (total, used) = ret[0]
            info.total = int(total)
            info.used = int(used)
        elif status == 'start':
            #zhenhao.he szl3bc12804 /dev/tty (v23) (szmaslic06/27003 337), start Thu 1/2 8:52
            pattern = '(\S+) (\S+).*\((\S+)/.*\), start (.*)'
            ret = re.findall(pattern, line)
            if not ret:
                continue
            used = UsedInfo()
            (used.user,used.host,used.lic,time_string) = ret[0]
            #处理跨年的情况
            year = datetime.now().year
            start_time = datetime.strptime(f"{year} {time_string}", "%Y %a %m/%d %H:%M")
            if start_time > datetime.now():
                year -= 1
                start_time = datetime.strptime(f"{year} {time_string}", "%Y %a %m/%d %H:%M")
            used.start_time = start_time
            # log.info(f"{lic} => {used}")
            info.used_info.append(used)
        # Next section
        elif line.startswith('Users of '):
            break
    return info

def draw_time_grpah(license_manager):
    info = license_manager.sorted()
    max_time = max(info.values())
    name_len = max([len(_) for _ in info.keys()])
    top = 100
    for user, time in info.items():
        curr = '█' * max(int(top * time / max_time), 1)
        name_remain = ' ' * (name_len - len(user))
        log.info(f"{user}{name_remain} :{time:5.2f}h {curr}")

if __name__ == '__main__':
    log.info("License status query...")
    lm = LicenseStatistic()
    log.info(lm)
    log.info(f"All licenses is: {[str(_) for _ in lm.get_lic_list()]}")
    for lic,number,total in lm.available():
        log.info(f"{lic}: {number} lic available")
    draw_time_grpah(lm)