Stanislav Hejny
98 lines
4.6 KiB
Python
98 lines
4.6 KiB
Python
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
|
|
|