"修改client的逻辑将queue的get阻塞式改成timout轮询的方式，以便于可以支持ctrc +c不过退出逻辑还存在问题需要修改。"
This commit is contained in:
parent
89e6c9439c
commit
26a8aecfb4
|
|
@ -10,8 +10,7 @@ import sys
|
|||
from my_log import get_logger
|
||||
from mock_lmutil import mock_license_add_user, mock_license_remove_user, _read_sample_file, _write_sample_file
|
||||
import queue
|
||||
import select # 添加 select 模块导入
|
||||
import signal
|
||||
import select
|
||||
|
||||
|
||||
log = get_logger()
|
||||
|
|
@ -31,25 +30,27 @@ class LicenseDispatcherClient:
|
|||
self.lic: queue.Queue = queue.Queue()
|
||||
self.current_lic: str = None
|
||||
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
|
||||
|
||||
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()
|
||||
if response['status'] == 'refuse':
|
||||
log.error(response['msg'])
|
||||
return None
|
||||
elif response['status'] == 'wait_dispatch':
|
||||
if response['status'] == 'wait_dispatch':
|
||||
# Start heartbeat thread
|
||||
log.info("Waiting for license allocation...")
|
||||
self._start_dispatcher_client(response['dispatcher_server']['port'])
|
||||
lic = self.get_lic()
|
||||
return lic
|
||||
else:
|
||||
log.warning(f"{response['status']}: {response['msg']}")
|
||||
return response['status']
|
||||
|
||||
def get_lic(self) -> str:
|
||||
"""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
|
||||
|
||||
def _request_license(self) -> dict:
|
||||
|
|
@ -71,7 +72,7 @@ class LicenseDispatcherClient:
|
|||
return json.loads(response_data)
|
||||
except (socket.timeout, ConnectionRefusedError) as e:
|
||||
log.error(f"Connection error: {e}")
|
||||
return {"status": "refuse", "msg": str(e)}
|
||||
return {"status": "connect_error", "msg": str(e)}
|
||||
except Exception as e:
|
||||
log.error(f"Requesting license error: {e}")
|
||||
return {"status": "error", "msg": str(e)}
|
||||
|
|
@ -83,7 +84,11 @@ class LicenseDispatcherClient:
|
|||
"""Start dispatcher connection with server"""
|
||||
log.info(f"Starting dispatcher: {port}")
|
||||
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()
|
||||
|
||||
def heartbeat_loop():
|
||||
|
|
@ -100,13 +105,13 @@ class LicenseDispatcherClient:
|
|||
|
||||
# Wait for response with timeout
|
||||
readable, _, _ = select.select([self.dispatcher_socket], [], [], 0.5)
|
||||
|
||||
if self.exit_event.is_set():
|
||||
break
|
||||
if self.dispatcher_socket in readable:
|
||||
data = self.dispatcher_socket.recv(1024)
|
||||
if not data:
|
||||
log.warning("Connection closed by server")
|
||||
break
|
||||
|
||||
response = json.loads(data.decode())
|
||||
if response.get('type') == 'lic_dispatch':
|
||||
self.lic.put(response['lic'])
|
||||
|
|
@ -121,12 +126,20 @@ class LicenseDispatcherClient:
|
|||
break
|
||||
elif response.get('type') == 'heartbeat_ack':
|
||||
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':
|
||||
wait_time = int(time.time() - start_time)
|
||||
log.info(f"waiting... {wait_time}s")
|
||||
|
||||
time.sleep(0.5) # Small sleep to prevent CPU spinning
|
||||
except KeyboardInterrupt:
|
||||
log.info("Keyboard interrupt received, shutting down...")
|
||||
break
|
||||
except Exception as e:
|
||||
log.error(f"Error in heartbeat loop: {e}")
|
||||
break
|
||||
|
|
@ -134,7 +147,7 @@ class LicenseDispatcherClient:
|
|||
self.lic.put(None) # Signal end of heartbeat loop
|
||||
|
||||
# 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()
|
||||
|
||||
def release(self) -> None:
|
||||
|
|
@ -150,37 +163,22 @@ class LicenseDispatcherClient:
|
|||
if self.heartbeat_thread and self.heartbeat_thread.is_alive():
|
||||
self.heartbeat_thread.join(timeout=1.0)
|
||||
|
||||
with self._lock:
|
||||
if self.dispatcher_socket:
|
||||
try:
|
||||
self.dispatcher_socket.shutdown(socket.SHUT_RDWR)
|
||||
except Exception:
|
||||
pass
|
||||
self.dispatcher_socket.close()
|
||||
self.dispatcher_socket = None
|
||||
if self.dispatcher_socket:
|
||||
try:
|
||||
self.dispatcher_socket.shutdown(socket.SHUT_RDWR)
|
||||
except Exception:
|
||||
pass
|
||||
self.dispatcher_socket.close()
|
||||
self.dispatcher_socket = None
|
||||
|
||||
self.release()
|
||||
self.lic.put(None)
|
||||
except Exception as e:
|
||||
log.error(f"Error during exit: {e}")
|
||||
|
||||
# Global flag for exit control
|
||||
is_exiting = False
|
||||
clients: dict[str, LicenseDispatcherClient] = {}
|
||||
|
||||
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"""
|
||||
def cleanup_clients():
|
||||
"""清理所有客户端资源"""
|
||||
for client_key, client in list(clients.items()):
|
||||
try:
|
||||
log.info(f"Cleaning up client: {client_key}")
|
||||
|
|
@ -189,43 +187,35 @@ def cleanup_and_exit():
|
|||
except Exception as e:
|
||||
log.error(f"Error cleaning up client {client_key}: {e}")
|
||||
log.info("\nProgram exited")
|
||||
sys.exit(0)
|
||||
|
||||
def main() -> None:
|
||||
# Register signal handlers
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
# Make sure SIGINT is not ignored
|
||||
if os.name == 'posix':
|
||||
signal.siginterrupt(signal.SIGINT, True)
|
||||
try:
|
||||
while True:
|
||||
show_menu()
|
||||
cmd = input("Please select an operation: ").strip()
|
||||
|
||||
while not is_exiting:
|
||||
show_menu()
|
||||
# Set timeout for input operations
|
||||
cmd = input("Please select an operation: ").strip()
|
||||
|
||||
if is_exiting: # Check if exit flag was set during input
|
||||
break
|
||||
|
||||
if cmd == '1':
|
||||
create_client()
|
||||
elif cmd == '2':
|
||||
release_client()
|
||||
elif cmd == '3':
|
||||
show_clients()
|
||||
elif cmd == '4':
|
||||
break
|
||||
else:
|
||||
print("Invalid command")
|
||||
if not is_exiting:
|
||||
cleanup_and_exit()
|
||||
if cmd == '1':
|
||||
create_client()
|
||||
elif cmd == '2':
|
||||
release_client()
|
||||
elif cmd == '3':
|
||||
show_clients()
|
||||
elif cmd == '4':
|
||||
break
|
||||
else:
|
||||
print("Invalid command")
|
||||
except KeyboardInterrupt:
|
||||
print("\nReceived keyboard interrupt, cleaning up...")
|
||||
finally:
|
||||
cleanup_clients()
|
||||
|
||||
def create_client():
|
||||
"""Create new license client"""
|
||||
global is_exiting
|
||||
username = input("Please enter username: ").strip()
|
||||
hostname = input("Please enter hostname: ").strip()
|
||||
if is_exiting:
|
||||
try:
|
||||
username = input("Please enter username: ").strip()
|
||||
hostname = input("Please enter hostname: ").strip()
|
||||
except KeyboardInterrupt:
|
||||
print("\nOperation cancelled")
|
||||
return
|
||||
|
||||
if not username or not hostname:
|
||||
|
|
@ -238,14 +228,17 @@ def create_client():
|
|||
return
|
||||
|
||||
client = LicenseDispatcherClient(user=username, host=hostname)
|
||||
lic = client.request()
|
||||
if lic:
|
||||
status = client.request()
|
||||
if status == 'wait_dispatch':
|
||||
lic = client.get_lic()
|
||||
clients[client_key] = client
|
||||
# Add mock license here
|
||||
mock_license_add_user(lic, username, hostname)
|
||||
print(f"{client_key} successfully created and requested license: {lic}")
|
||||
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():
|
||||
"""Release existing license"""
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user