feat: add Consul-based global ID generator with fallback mechanism

amqp/service/consul_global_id_generator.py

@@ -0,0 +1,97 @@
+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
+