#!/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