Stanislav Hejny
98 lines
3.8 KiB
Python
98 lines
3.8 KiB
Python
import logging
|
|
import os
|
|
import random
|
|
import socket
|
|
import time
|
|
|
|
import consul_kv
|
|
|
|
from amqp.config.amq_configuration import AMQAdapter
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_config_values(config: AMQAdapter) -> tuple:
|
|
"""
|
|
Get configuration values from AMQConfiguration or environment variables.
|
|
|
|
Returns:
|
|
tuple: (consul_host, consul_port, consul_counter_key, max_retries, retry_delay_seconds)
|
|
"""
|
|
try:
|
|
consul_host = config.consul_host
|
|
consul_port = config.consul_port
|
|
consul_counter_key = config.consul_counter_key
|
|
max_retries = config.consul_max_retries
|
|
retry_delay_seconds = config.consul_initial_retry_delay
|
|
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", consul_kv.settings.DEFAULT_ENDPOINT.split(":")[0]),
|
|
int(os.environ.get("CONSUL_PORT", consul_kv.settings.DEFAULT_ENDPOINT.split(":")[1])),
|
|
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.2)),
|
|
)
|
|
|
|
|
|
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
|
|
"""
|
|
timestamp = int(time.time() * 1000)
|
|
hostname = socket.gethostname()
|
|
hostname_hash = hash(hostname) % 10000
|
|
random_part = random.randint(0, 9999)
|
|
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(config: AMQAdapter) -> 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(config)
|
|
)
|
|
try:
|
|
c = consul_kv.Connection(endpoint=f"{consul_host}:{consul_port}", timeout=5)
|
|
index, data = c.get(consul_counter_key)
|
|
if data is None:
|
|
logger.info(f"Initializing Consul counter at {consul_counter_key}")
|
|
if c.put(consul_counter_key, "1", cas=1):
|
|
return 1
|
|
else:
|
|
logger.error("Failed to initialize Consul counter")
|
|
return generate_fallback_id()
|
|
current_value = int(data["Value"].decode("utf-8"))
|
|
new_value = current_value + 1
|
|
for attempt in range(max_retries):
|
|
success = c.put(consul_counter_key, str(new_value), cas=data["ModifyIndex"])
|
|
if success:
|
|
logger.debug(f"Successfully obtained unique ID: {new_value}")
|
|
return new_value
|
|
logger.debug(f"CAS update failed on attempt {attempt + 1}, retrying...")
|
|
time.sleep(retry_delay_seconds * (1 << attempt)) # Exponential backoff
|
|
index, data = c.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
|
|
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()
|