amq-adapter-python/amqp/service/consul_global_id_generator.py

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
    """
    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()