import consul_kv import socket import os import time import logging import threading # --- Configuration --- CONSUL_HOST = os.environ.get('CONSUL_HOST', 'localhost') CONSUL_PORT = int(os.environ.get('CONSUL_PORT', 8500)) CONSUL_COUNTER_KEY = 'service/ids/counter' MAX_RETRIES = 5 # Max attempts for atomic update RETRY_DELAY_SECONDS = 0.1 # Delay between retries for atomic update # --- Logging Setup --- logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(threadName)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # --- Fallback ID Generation --- def generate_fallback_id() -> int: """ Generates a fallback ID based on hostname and PID if Consul is unavailable. This provides a unique-ish ID for local testing/development but is not guaranteed to be globally unique across different machines or after restarts. """ hostname = socket.gethostname() pid = os.getpid() # A simple way to combine them into a number. # Hash hostname to an integer and combine with PID. # Note: This is a simple fallback. For production, consider a more robust # local ID generation or a different strategy if Consul is critical. fallback_id = (hash(hostname) % 1000000) * 100000 + pid logger.warning( f"Consul is unavailable. Generating fallback ID: {fallback_id} (based on hostname: {hostname}, PID: {pid})") return fallback_id # --- Unique Instance ID Generation Function --- def get_unique_instance_id() -> int: """ Reads and atomically increments a counter in Consul KV to get a unique instance ID. If Consul is unavailable or the atomic update fails after retries, it falls back to generating an ID based on hostname and PID. """ try: # Initialize Consul client client = consul_kv.Connection(endpoint=f"{CONSUL_HOST}:{CONSUL_PORT}") logger.info(f"Attempting to connect to Consul at {CONSUL_HOST}:{CONSUL_PORT}") for attempt in range(MAX_RETRIES): try: # 1. Read the current value and its modifyIndex (for CAS) # client.get returns (value, modifyIndex) current_value_str, modify_index = client.get(CONSUL_COUNTER_KEY) # If key doesn't exist, initialize it to 0 if current_value_str is None: current_value = 0 modify_index = 0 # For a non-existent key, modify_index is 0 for initial CAS logger.info(f"Consul key '{CONSUL_COUNTER_KEY}' not found. Initializing to 0.") else: current_value = int(current_value_str) new_value = current_value + 1 # 2. Attempt atomic update using CAS (Check-And-Set) # client.set returns True on success, False on failure (if modify_index changed) success = client.put(CONSUL_COUNTER_KEY, str(new_value), cas=modify_index) if success: logger.info(f"Successfully obtained unique ID: {new_value} (Attempt {attempt + 1}/{MAX_RETRIES})") return new_value else: logger.warning( f"CAS failed for '{CONSUL_COUNTER_KEY}' (modify_index changed). Retrying... (Attempt {attempt + 1}/{MAX_RETRIES})") time.sleep(RETRY_DELAY_SECONDS * (attempt + 1)) # Exponential backoff for retries except ValueError: # Handle case where the stored value is not an integer logger.error( f"Value stored at '{CONSUL_COUNTER_KEY}' is not a valid integer. Please check Consul KV. Attempting to overwrite.") # To recover, you might force a set without CAS, but that risks data loss. # For this demo, we'll just let it retry or fall back. time.sleep(RETRY_DELAY_SECONDS * (attempt + 1)) except Exception as e: # Catch specific ConsulKV errors (e.g., connection issues, invalid response) logger.error(f"ConsulKV error during atomic update (Attempt {attempt + 1}/{MAX_RETRIES}): {e}") time.sleep(RETRY_DELAY_SECONDS * (attempt + 1)) # Exponential backoff logger.error(f"Failed to obtain unique ID from Consul after {MAX_RETRIES} attempts.") return generate_fallback_id() # Fallback if all retries fail except Exception as e: # Catch broader connection errors or unexpected issues with Consul client initialization logger.error(f"Could not connect to Consul or an unexpected error occurred: {e}") return generate_fallback_id() # Fallback if Consul is completely unreachable import os import logging import random import socket import time from typing import Optional import consul from amqp.config.amq_configuration import AMQConfiguration from amqp.config.amq_configuration import AMQConfiguration # Configuration from environment variables with defaults CONSUL_HOST = os.environ.get('CONSUL_HOST', 'localhost') CONSUL_PORT = int(os.environ.get('CONSUL_PORT', 8500)) CONSUL_COUNTER_KEY = 'service/ids/counter' MAX_RETRIES = 5 # Max attempts for atomic update RETRY_DELAY_SECONDS = 0.1 # Delay between retries for atomic update logger = logging.getLogger(__name__) def generate_fallback_id() -> int: """ Generate a fallback ID when Consul is not available. Uses a combination of timestamp, hostname hash, and random number. Returns: int: A reasonably unique integer ID """ # Get current timestamp in milliseconds timestamp = int(time.time() * 1000) # Get hostname and hash it to an integer hostname = socket.gethostname() hostname_hash = hash(hostname) % 10000 # Limit to 4 digits # Generate a random number random_part = random.randint(0, 9999) # 4 digits # Combine all parts into a single integer # Format: TTTTTTTTTTT-HHHH-RRRR (T=timestamp, H=hostname hash, R=random) unique_id = (timestamp * 100000000) + (hostname_hash * 10000) + random_part logger.warning(f"Using fallback ID generation method: {unique_id}") return unique_id def get_unique_instance_id() -> int: """ Get a globally unique ID from Consul. Uses Consul's atomic Compare-And-Set operations to safely increment a counter. Falls back to a local generation method if Consul is unavailable. Returns: int: A globally unique integer ID """ consul_host, consul_port, consul_counter_key, max_retries, retry_delay_seconds = get_config_values() try: # Connect to Consul c = consul.Consul(host=consul_host, port=consul_port) # Try to get the current counter value index, data = c.kv.get(CONSUL_COUNTER_KEY) # If the key doesn't exist yet, create it with initial value if data is None: logger.info(f"Initializing Consul counter at {CONSUL_COUNTER_KEY}") if c.kv.put(CONSUL_COUNTER_KEY, "1"): return 1 else: logger.error("Failed to initialize Consul counter") return generate_fallback_id() # Get the current value current_value = int(data['Value'].decode('utf-8')) new_value = current_value + 1 # Try to atomically update the counter for attempt in range(MAX_RETRIES): # Use Compare-And-Set to ensure atomicity success = c.kv.put( CONSUL_COUNTER_KEY, str(new_value), cas=data['ModifyIndex'] ) if success: logger.debug(f"Successfully obtained unique ID: {new_value}") return new_value # If CAS failed, someone else updated the value, retry logger.debug(f"CAS update failed on attempt {attempt+1}, retrying...") time.sleep(RETRY_DELAY_SECONDS) # Get the latest value for the next attempt index, data = c.kv.get(CONSUL_COUNTER_KEY) if data is None: logger.error("Counter disappeared during update") return generate_fallback_id() current_value = int(data['Value'].decode('utf-8')) new_value = current_value + 1 # If we've exhausted all retries, use the fallback logger.error(f"Failed to update counter after {MAX_RETRIES} attempts") return generate_fallback_id() except Exception as e: logger.error(f"Error accessing Consul: {str(e)}") return generate_fallback_id() def get_config_values(): """ Get configuration values from AMQConfiguration or environment variables. Returns: tuple: (consul_host, consul_port, consul_counter_key, max_retries, retry_delay_seconds) """ try: config = AMQConfiguration("application.properties") consul_host = os.environ.get('CONSUL_HOST', 'localhost') consul_port = int(os.environ.get('CONSUL_PORT', 8500)) consul_counter_key = os.environ.get('CONSUL_COUNTER_KEY', 'service/ids/counter') max_retries = int(os.environ.get('CONSUL_MAX_RETRIES', 5)) retry_delay_seconds = float(os.environ.get('CONSUL_RETRY_DELAY_SECONDS', 0.1)) return consul_host, consul_port, consul_counter_key, max_retries, retry_delay_seconds except Exception as e: logger.warning(f"Failed to load configuration: {str(e)}. Using default values.") return ( os.environ.get('CONSUL_HOST', 'localhost'), int(os.environ.get('CONSUL_PORT', 8500)), os.environ.get('CONSUL_COUNTER_KEY', 'service/ids/counter'), int(os.environ.get('CONSUL_MAX_RETRIES', 5)), float(os.environ.get('CONSUL_RETRY_DELAY_SECONDS', 0.1)) )