ekko/vcast_license_usage_server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

"修改client的逻辑将queue的get阻塞式改成timout轮询的方式，以便于可以支持ctrc +c不过退出逻辑还存在问题需要修改。"

Browse Source
This commit is contained in:
Ekko.bao 2025-01-14 09:01:00 +08:00
parent 89e6c9439c
commit 26a8aecfb4
1 changed files with 69 additions and 76 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

145
license_client.py
View File

@ -10,8 +10,7 @@ import sys
from my_log import get_logger from my_log import get_logger
from mock_lmutil import mock_license_add_user, mock_license_remove_user, _read_sample_file, _write_sample_file from mock_lmutil import mock_license_add_user, mock_license_remove_user, _read_sample_file, _write_sample_file
import queue import queue
import select # 添加 select 模块导入 import select
import signal
log = get_logger() log = get_logger()
@ -31,25 +30,27 @@ class LicenseDispatcherClient:
self.lic: queue.Queue = queue.Queue() self.lic: queue.Queue = queue.Queue()
self.current_lic: str = None self.current_lic: str = None
self.dispatcher_timer: threading.Timer = None # Add missing attribute self.dispatcher_timer: threading.Timer = None # Add missing attribute
self._lock = threading.Lock() # Add thread safety
self.heartbeat_thread: threading.Thread | None = None # Add thread reference self.heartbeat_thread: threading.Thread | None = None # Add thread reference
def request(self) -> str: def request(self) -> str:
"""Request license, block until license is obtained or user interrupts""" """Request license, block until license is obtained"""
response = self._request_license() response = self._request_license()
if response['status'] == 'refuse': if response['status'] == 'wait_dispatch':
log.error(response['msg'])
return None
elif response['status'] == 'wait_dispatch':
# Start heartbeat thread # Start heartbeat thread
log.info("Waiting for license allocation...") log.info("Waiting for license allocation...")
self._start_dispatcher_client(response['dispatcher_server']['port']) self._start_dispatcher_client(response['dispatcher_server']['port'])
lic = self.get_lic() else:
return lic log.warning(f"{response['status']}: {response['msg']}")
return response['status']
def get_lic(self) -> str: def get_lic(self) -> str:
"""Get the assigned license""" """Get the assigned license"""
self.current_lic = self.lic.get() while not self.exit_event.is_set():
try:
self.current_lic = self.lic.get(timeout=0.1)
return self.current_lic
except queue.Empty:
continue
return self.current_lic return self.current_lic
def _request_license(self) -> dict: def _request_license(self) -> dict:
@ -71,7 +72,7 @@ class LicenseDispatcherClient:
return json.loads(response_data) return json.loads(response_data)
except (socket.timeout, ConnectionRefusedError) as e: except (socket.timeout, ConnectionRefusedError) as e:
log.error(f"Connection error: {e}") log.error(f"Connection error: {e}")
return {"status": "refuse", "msg": str(e)} return {"status": "connect_error", "msg": str(e)}
except Exception as e: except Exception as e:
log.error(f"Requesting license error: {e}") log.error(f"Requesting license error: {e}")
return {"status": "error", "msg": str(e)} return {"status": "error", "msg": str(e)}
@ -83,7 +84,11 @@ class LicenseDispatcherClient:
"""Start dispatcher connection with server""" """Start dispatcher connection with server"""
log.info(f"Starting dispatcher: {port}") log.info(f"Starting dispatcher: {port}")
self.dispatcher_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.dispatcher_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.dispatcher_socket.connect((SERVER_HOST, port)) try:
self.dispatcher_socket.connect((SERVER_HOST, port))
except Exception as e:
log.error(f"Failed to connect to dispatcher: {e}")
return
start_time = time.time() start_time = time.time()
def heartbeat_loop(): def heartbeat_loop():
@ -100,13 +105,13 @@ class LicenseDispatcherClient:
# Wait for response with timeout # Wait for response with timeout
readable, _, _ = select.select([self.dispatcher_socket], [], [], 0.5) readable, _, _ = select.select([self.dispatcher_socket], [], [], 0.5)
if self.exit_event.is_set():
break
if self.dispatcher_socket in readable: if self.dispatcher_socket in readable:
data = self.dispatcher_socket.recv(1024) data = self.dispatcher_socket.recv(1024)
if not data: if not data:
log.warning("Connection closed by server") log.warning("Connection closed by server")
break break
response = json.loads(data.decode()) response = json.loads(data.decode())
if response.get('type') == 'lic_dispatch': if response.get('type') == 'lic_dispatch':
self.lic.put(response['lic']) self.lic.put(response['lic'])
@ -121,12 +126,20 @@ class LicenseDispatcherClient:
break break
elif response.get('type') == 'heartbeat_ack': elif response.get('type') == 'heartbeat_ack':
wait_time = int(time.time() - start_time) wait_time = int(time.time() - start_time)
print(f"\rwaiting... {wait_time}s", end='', flush=True) if wait_time < 60:
print(f"\rwaiting... {wait_time}s", end='', flush=True)
elif wait_time < 3600:
print(f"\rwaiting... {wait_time / 60}m {wait_time}s", end='', flush=True)
else:
print(f"\rwaiting... {wait_time / 3600}h {wait_time / 60}m {wait_time}s", end='', flush=True)
elif response.get('type') == 'queue': elif response.get('type') == 'queue':
wait_time = int(time.time() - start_time) wait_time = int(time.time() - start_time)
log.info(f"waiting... {wait_time}s") log.info(f"waiting... {wait_time}s")
time.sleep(0.5) # Small sleep to prevent CPU spinning time.sleep(0.5) # Small sleep to prevent CPU spinning
except KeyboardInterrupt:
log.info("Keyboard interrupt received, shutting down...")
break
except Exception as e: except Exception as e:
log.error(f"Error in heartbeat loop: {e}") log.error(f"Error in heartbeat loop: {e}")
break break
@ -134,7 +147,7 @@ class LicenseDispatcherClient:
self.lic.put(None) # Signal end of heartbeat loop self.lic.put(None) # Signal end of heartbeat loop
# Start heartbeat thread # Start heartbeat thread
self.heartbeat_thread = threading.Thread(target=heartbeat_loop, daemon=True) self.heartbeat_thread = threading.Thread(target=heartbeat_loop)
self.heartbeat_thread.start() self.heartbeat_thread.start()
def release(self) -> None: def release(self) -> None:
@ -150,37 +163,22 @@ class LicenseDispatcherClient:
if self.heartbeat_thread and self.heartbeat_thread.is_alive(): if self.heartbeat_thread and self.heartbeat_thread.is_alive():
self.heartbeat_thread.join(timeout=1.0) self.heartbeat_thread.join(timeout=1.0)
with self._lock: if self.dispatcher_socket:
if self.dispatcher_socket: try:
try: self.dispatcher_socket.shutdown(socket.SHUT_RDWR)
self.dispatcher_socket.shutdown(socket.SHUT_RDWR) except Exception:
except Exception: pass
pass self.dispatcher_socket.close()
self.dispatcher_socket.close() self.dispatcher_socket = None
self.dispatcher_socket = None
self.release() self.release()
self.lic.put(None) self.lic.put(None)
except Exception as e: except Exception as e:
log.error(f"Error during exit: {e}") log.error(f"Error during exit: {e}")
# Global flag for exit control
is_exiting = False
clients: dict[str, LicenseDispatcherClient] = {} clients: dict[str, LicenseDispatcherClient] = {}
def cleanup_clients():
def signal_handler(signum, frame): """清理所有客户端资源"""
"""Handle interrupt signals"""
global is_exiting
if is_exiting:
log.info("\nForce exiting...")
sys.exit(1)
log.info("\nReceived interrupt signal, cleaning up...")
is_exiting = True
cleanup_and_exit()
def cleanup_and_exit():
"""Cleanup all resources and exit"""
for client_key, client in list(clients.items()): for client_key, client in list(clients.items()):
try: try:
log.info(f"Cleaning up client: {client_key}") log.info(f"Cleaning up client: {client_key}")
@ -189,43 +187,35 @@ def cleanup_and_exit():
except Exception as e: except Exception as e:
log.error(f"Error cleaning up client {client_key}: {e}") log.error(f"Error cleaning up client {client_key}: {e}")
log.info("\nProgram exited") log.info("\nProgram exited")
sys.exit(0)
def main() -> None: def main() -> None:
# Register signal handlers try:
signal.signal(signal.SIGINT, signal_handler) while True:
signal.signal(signal.SIGTERM, signal_handler) show_menu()
# Make sure SIGINT is not ignored cmd = input("Please select an operation: ").strip()
if os.name == 'posix':
signal.siginterrupt(signal.SIGINT, True)
while not is_exiting: if cmd == '1':
show_menu() create_client()
# Set timeout for input operations elif cmd == '2':
cmd = input("Please select an operation: ").strip() release_client()
elif cmd == '3':
if is_exiting: # Check if exit flag was set during input show_clients()
break elif cmd == '4':
break
if cmd == '1': else:
create_client() print("Invalid command")
elif cmd == '2': except KeyboardInterrupt:
release_client() print("\nReceived keyboard interrupt, cleaning up...")
elif cmd == '3': finally:
show_clients() cleanup_clients()
elif cmd == '4':
break
else:
print("Invalid command")
if not is_exiting:
cleanup_and_exit()
def create_client(): def create_client():
"""Create new license client""" """Create new license client"""
global is_exiting try:
username = input("Please enter username: ").strip() username = input("Please enter username: ").strip()
hostname = input("Please enter hostname: ").strip() hostname = input("Please enter hostname: ").strip()
if is_exiting: except KeyboardInterrupt:
print("\nOperation cancelled")
return return
if not username or not hostname: if not username or not hostname:
@ -238,14 +228,17 @@ def create_client():
return return
client = LicenseDispatcherClient(user=username, host=hostname) client = LicenseDispatcherClient(user=username, host=hostname)
lic = client.request() status = client.request()
if lic: if status == 'wait_dispatch':
lic = client.get_lic()
clients[client_key] = client clients[client_key] = client
# Add mock license here
mock_license_add_user(lic, username, hostname) mock_license_add_user(lic, username, hostname)
print(f"{client_key} successfully created and requested license: {lic}") print(f"{client_key} successfully created and requested license: {lic}")
else: else:
print(f"{client_key} License request failed") if status == 'connect_error':
print(f"{client_key} License Server connect error, You can try again later")
else:
print(f"{client_key} License request failed: {status}")
def release_client(): def release_client():
"""Release existing license""" """Release existing license"""