148 lines
6.4 KiB
Python
148 lines
6.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
CLI tool for acquiring and releasing profile locks.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
try:
|
|
from dotenv import load_dotenv
|
|
except ImportError:
|
|
load_dotenv = None
|
|
|
|
from .profile_manager_tool import ProfileManager
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def add_profile_allocator_parser(subparsers):
|
|
"""Adds the parser for the 'profile-allocator' command."""
|
|
parser = subparsers.add_parser(
|
|
'profile-allocator',
|
|
description='Acquire and release profile locks.',
|
|
formatter_class=argparse.RawTextHelpFormatter,
|
|
help='Acquire and release profile locks.'
|
|
)
|
|
|
|
common_parser = argparse.ArgumentParser(add_help=False)
|
|
common_parser.add_argument('--env-file', help='Path to a .env file to load environment variables from.')
|
|
common_parser.add_argument('--redis-host', default=None, help='Redis host. Defaults to MASTER_HOST_IP or REDIS_HOST env var, or localhost.')
|
|
common_parser.add_argument('--redis-port', type=int, default=None, help='Redis port. Defaults to REDIS_PORT env var, or 6379.')
|
|
common_parser.add_argument('--redis-password', default=None, help='Redis password. Defaults to REDIS_PASSWORD env var.')
|
|
common_parser.add_argument('--env', default='dev', help="Environment name for Redis key prefix (e.g., 'stg', 'prod'). Defaults to 'dev'.")
|
|
common_parser.add_argument('--legacy', action='store_true', help="Use legacy key prefix ('profile_mgmt_') without environment.")
|
|
common_parser.add_argument('--key-prefix', default=None, help='Explicit key prefix for Redis. Overrides --env, --legacy and any defaults.')
|
|
common_parser.add_argument('--verbose', action='store_true', help='Enable verbose logging')
|
|
|
|
allocator_subparsers = parser.add_subparsers(dest='allocator_command', help='Command to execute', required=True)
|
|
|
|
# Lock command
|
|
lock_parser = allocator_subparsers.add_parser('lock', help='Find and lock an available profile', parents=[common_parser])
|
|
lock_parser.add_argument('--owner', required=True, help='Identifier for the process locking the profile')
|
|
lock_parser.add_argument('--profile-prefix', help='Only lock profiles with this name prefix')
|
|
lock_parser.add_argument('--wait', action='store_true', help='Wait indefinitely for a profile to become available, with exponential backoff.')
|
|
|
|
# Unlock command
|
|
unlock_parser = allocator_subparsers.add_parser('unlock', help='Unlock a profile', parents=[common_parser])
|
|
unlock_parser.add_argument('name', help='Profile name to unlock')
|
|
unlock_parser.add_argument('--owner', help='Identifier of the owner. If provided, unlock will only succeed if owner matches.')
|
|
|
|
# Cleanup command
|
|
cleanup_parser = allocator_subparsers.add_parser('cleanup-locks', help='Clean up stale locks', parents=[common_parser])
|
|
cleanup_parser.add_argument('--max-age-seconds', type=int, default=3600,
|
|
help='Maximum lock age in seconds before it is considered stale (default: 3600)')
|
|
|
|
return parser
|
|
|
|
def main_profile_allocator(args):
|
|
"""Main dispatcher for 'profile-allocator' command."""
|
|
if load_dotenv:
|
|
env_file = args.env_file
|
|
if not env_file and args.env and '.env' in args.env and os.path.exists(args.env):
|
|
logger.warning(f"Warning: --env should be an environment name (e.g., 'dev'), not a file path. Treating '{args.env}' as --env-file. The environment name will default to 'dev'.")
|
|
env_file = args.env
|
|
args.env = 'dev'
|
|
|
|
was_loaded = load_dotenv(env_file)
|
|
if was_loaded:
|
|
logger.info(f"Loaded environment variables from {env_file or '.env file'}")
|
|
elif args.env_file:
|
|
logger.error(f"The specified --env-file was not found: {args.env_file}")
|
|
return 1
|
|
|
|
if args.redis_host is None:
|
|
args.redis_host = os.getenv('MASTER_HOST_IP', os.getenv('REDIS_HOST', 'localhost'))
|
|
if args.redis_port is None:
|
|
args.redis_port = int(os.getenv('REDIS_PORT', 6379))
|
|
if args.redis_password is None:
|
|
args.redis_password = os.getenv('REDIS_PASSWORD')
|
|
|
|
if args.verbose:
|
|
logging.getLogger().setLevel(logging.DEBUG)
|
|
|
|
if args.key_prefix:
|
|
key_prefix = args.key_prefix
|
|
elif args.legacy:
|
|
key_prefix = 'profile_mgmt_'
|
|
else:
|
|
key_prefix = f"{args.env}_profile_mgmt_"
|
|
|
|
manager = ProfileManager(
|
|
redis_host=args.redis_host,
|
|
redis_port=args.redis_port,
|
|
redis_password=args.redis_password,
|
|
key_prefix=key_prefix
|
|
)
|
|
|
|
if args.allocator_command == 'lock':
|
|
if not args.wait:
|
|
profile = manager.lock_profile(args.owner, profile_prefix=args.profile_prefix)
|
|
if profile:
|
|
print(json.dumps(profile, indent=2, default=str))
|
|
return 0
|
|
else:
|
|
print("No available profile could be locked.", file=sys.stderr)
|
|
return 1
|
|
|
|
# With --wait, loop with backoff
|
|
lock_attempts = 0
|
|
backoff_seconds = [3, 5, 9, 20, 50, 120, 300]
|
|
while True:
|
|
profile = manager.lock_profile(args.owner, profile_prefix=args.profile_prefix)
|
|
if profile:
|
|
print(json.dumps(profile, indent=2, default=str))
|
|
return 0
|
|
|
|
sleep_duration = backoff_seconds[min(lock_attempts, len(backoff_seconds) - 1)]
|
|
logger.info(f"No available profile. Retrying in {sleep_duration}s... (attempt {lock_attempts + 1})")
|
|
|
|
try:
|
|
time.sleep(sleep_duration)
|
|
except KeyboardInterrupt:
|
|
logger.warning("Wait for lock interrupted by user.")
|
|
# Use print for stderr as well, since logger might be configured differently by callers
|
|
print("\nWait for lock interrupted by user.", file=sys.stderr)
|
|
return 130 # Standard exit code for Ctrl+C
|
|
|
|
lock_attempts += 1
|
|
|
|
elif args.allocator_command == 'unlock':
|
|
success = manager.unlock_profile(args.name, args.owner)
|
|
return 0 if success else 1
|
|
|
|
elif args.allocator_command == 'cleanup-locks':
|
|
cleaned_count = manager.cleanup_stale_locks(args.max_age_seconds)
|
|
print(f"Cleaned {cleaned_count} stale lock(s).")
|
|
return 0
|
|
|
|
return 1 # Should not be reached
|