diff --git a/amqp/service/consul_global_id_generator.py b/amqp/service/consul_global_id_generator.py index bb4e988..9c693c4 100644 --- a/amqp/service/consul_global_id_generator.py +++ b/amqp/service/consul_global_id_generator.py @@ -95,3 +95,110 @@ def get_unique_instance_id() -> int: 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 + +# 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 + """ + 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()