Initial commit
This commit is contained in:
285
scripts/worker.py
Normal file
285
scripts/worker.py
Normal file
@@ -0,0 +1,285 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
NATS Worker that consumes messages and routes them to different external endpoints
|
||||
Based on the endpoint in the message, forwards to the corresponding external API
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import asyncio
|
||||
import aiohttp
|
||||
import nats
|
||||
from prometheus_client import Counter, Histogram, Gauge, start_http_server
|
||||
import signal
|
||||
import sys
|
||||
|
||||
# Prometheus metrics
|
||||
messages_processed = Counter('worker_messages_processed_total', 'Total messages processed', ['status', 'endpoint'])
|
||||
message_duration = Histogram('worker_message_duration_seconds', 'Message processing duration', ['endpoint'])
|
||||
messages_in_flight = Gauge('worker_messages_in_flight', 'Messages currently being processed')
|
||||
queue_depth = Gauge('worker_queue_depth', 'Approximate queue depth')
|
||||
|
||||
# Configuration
|
||||
NATS_URL = os.getenv("NATS_URL", "nats://nats-server:4222")
|
||||
NATS_USER = os.getenv("NATS_USER", "admin")
|
||||
NATS_PASSWORD = os.getenv("NATS_PASSWORD", "package@321#")
|
||||
NATS_STREAM = os.getenv("NATS_STREAM", "EVENTS")
|
||||
NATS_SUBJECT = os.getenv("NATS_SUBJECT", "api.>") # Subscribe to all API subjects
|
||||
NATS_CONSUMER = os.getenv("NATS_CONSUMER", "worker_consumer")
|
||||
WORKER_CONCURRENCY = int(os.getenv("WORKER_CONCURRENCY", "10"))
|
||||
RETRY_ATTEMPTS = int(os.getenv("RETRY_ATTEMPTS", "5"))
|
||||
RETRY_DELAY = int(os.getenv("RETRY_DELAY_SECONDS", "5"))
|
||||
|
||||
# Base URL for external APIs
|
||||
BASE_URL = os.getenv("EXTERNAL_BASE_URL", "https://jupiter.nearle.app")
|
||||
|
||||
# Endpoint mapping: FastAPI endpoint -> External API endpoint
|
||||
ENDPOINT_MAPPING = {
|
||||
"/live/api/v1/deliveries/createdeliveries": f"{BASE_URL}/live/api/v1/deliveries/createdeliveries",
|
||||
"/live/api/v1/deliveries/updatedelivery": f"{BASE_URL}/live/api/v1/deliveries/updatedelivery",
|
||||
"/live/api/v2/partners/createriderlog": f"{BASE_URL}/live/api/v2/partners/createriderlog",
|
||||
"/live/api/v2/deliveries/createdeliverylog": f"{BASE_URL}/live/api/v2/deliveries/createdeliverylog",
|
||||
"/live/api/v2/partners/createbreaklog": f"{BASE_URL}/live/api/v2/partners/createbreaklog",
|
||||
"/live/api/v2/partners/updatebreaklog": f"{BASE_URL}/live/api/v2/partners/updatebreaklog",
|
||||
}
|
||||
|
||||
# Global state
|
||||
nc = None
|
||||
js = None
|
||||
running = True
|
||||
semaphore = asyncio.Semaphore(WORKER_CONCURRENCY)
|
||||
|
||||
def signal_handler(sig, frame):
|
||||
"""Handle shutdown signals"""
|
||||
global running
|
||||
print("\n🛑 Shutdown signal received, stopping worker...")
|
||||
running = False
|
||||
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
|
||||
async def forward_to_external(endpoint: str, payload: dict, api_key: str = None, method: str = "POST") -> bool:
|
||||
"""Forward message to external endpoint with retries"""
|
||||
external_url = ENDPOINT_MAPPING.get(endpoint)
|
||||
|
||||
if not external_url:
|
||||
print(f"⚠️ Unknown endpoint: {endpoint}, skipping...")
|
||||
return False
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
|
||||
# Get the actual data to forward
|
||||
data_to_forward = payload.get("data", payload)
|
||||
|
||||
# Normalize payload shape for endpoints that expect an array
|
||||
if endpoint == "/live/api/v2/deliveries/createdeliverylog":
|
||||
if isinstance(data_to_forward, dict):
|
||||
data_to_forward = [data_to_forward]
|
||||
elif not isinstance(data_to_forward, list):
|
||||
# Fallback to wrapping anything that's not already list-like
|
||||
data_to_forward = [data_to_forward]
|
||||
|
||||
# Debug logging for delivery log payload shape
|
||||
print(f"➡️ Forwarding createdeliverylog payload: {data_to_forward}")
|
||||
|
||||
# Normalize method to uppercase
|
||||
http_method = method.upper() if method else "POST"
|
||||
|
||||
for attempt in range(RETRY_ATTEMPTS):
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Use the correct HTTP method
|
||||
if http_method == "GET":
|
||||
async with session.get(
|
||||
external_url,
|
||||
params=data_to_forward if isinstance(data_to_forward, dict) else {},
|
||||
headers=headers,
|
||||
timeout=aiohttp.ClientTimeout(total=30)
|
||||
) as response:
|
||||
return await _handle_response(response, endpoint, attempt)
|
||||
elif http_method == "PUT":
|
||||
async with session.put(
|
||||
external_url,
|
||||
json=data_to_forward,
|
||||
headers=headers,
|
||||
timeout=aiohttp.ClientTimeout(total=30)
|
||||
) as response:
|
||||
return await _handle_response(response, endpoint, attempt)
|
||||
elif http_method == "PATCH":
|
||||
async with session.patch(
|
||||
external_url,
|
||||
json=data_to_forward,
|
||||
headers=headers,
|
||||
timeout=aiohttp.ClientTimeout(total=30)
|
||||
) as response:
|
||||
return await _handle_response(response, endpoint, attempt)
|
||||
elif http_method == "DELETE":
|
||||
async with session.delete(
|
||||
external_url,
|
||||
headers=headers,
|
||||
timeout=aiohttp.ClientTimeout(total=30)
|
||||
) as response:
|
||||
return await _handle_response(response, endpoint, attempt)
|
||||
else: # Default to POST
|
||||
async with session.post(
|
||||
external_url,
|
||||
json=data_to_forward,
|
||||
headers=headers,
|
||||
timeout=aiohttp.ClientTimeout(total=30)
|
||||
) as response:
|
||||
return await _handle_response(response, endpoint, attempt)
|
||||
except asyncio.TimeoutError:
|
||||
print(f"⚠️ Timeout forwarding {endpoint} to external API (attempt {attempt + 1}/{RETRY_ATTEMPTS})")
|
||||
if attempt < RETRY_ATTEMPTS - 1:
|
||||
await asyncio.sleep(RETRY_DELAY * (attempt + 1))
|
||||
except Exception as e:
|
||||
print(f"⚠️ Error forwarding {endpoint} to external API: {e} (attempt {attempt + 1}/{RETRY_ATTEMPTS})")
|
||||
if attempt < RETRY_ATTEMPTS - 1:
|
||||
await asyncio.sleep(RETRY_DELAY * (attempt + 1))
|
||||
|
||||
return False
|
||||
|
||||
async def _handle_response(response, endpoint: str, attempt: int) -> bool:
|
||||
"""Handle HTTP response and determine if we should retry"""
|
||||
body_text = None
|
||||
try:
|
||||
body_text = await response.text()
|
||||
except Exception:
|
||||
body_text = "<no body>"
|
||||
|
||||
if response.status == 200 or response.status == 201:
|
||||
return True
|
||||
elif response.status >= 500:
|
||||
# Server error - retry
|
||||
print(f"⚠️ External API returned {response.status} for {endpoint}, retrying... (attempt {attempt + 1}/{RETRY_ATTEMPTS}) body={body_text}")
|
||||
if attempt < RETRY_ATTEMPTS - 1:
|
||||
await asyncio.sleep(RETRY_DELAY * (attempt + 1)) # Exponential backoff
|
||||
return False
|
||||
else:
|
||||
# Client error - don't retry
|
||||
print(f"❌ External API returned {response.status} for {endpoint}, not retrying; body={body_text}")
|
||||
return False
|
||||
|
||||
async def process_message(msg):
|
||||
"""Process a single message"""
|
||||
async with semaphore:
|
||||
messages_in_flight.inc()
|
||||
start_time = asyncio.get_event_loop().time()
|
||||
endpoint = "unknown"
|
||||
|
||||
try:
|
||||
# Decode message
|
||||
payload = json.loads(msg.data.decode())
|
||||
endpoint = payload.get("endpoint", payload.get("original_path", "unknown"))
|
||||
http_method = payload.get("method", "POST") # Get HTTP method from payload
|
||||
|
||||
print(f"📨 Processing message for endpoint: {endpoint} (method: {http_method})")
|
||||
|
||||
# Forward to external endpoint with correct HTTP method
|
||||
api_key = os.getenv("EXTERNAL_ENDPOINT_API_KEY", "")
|
||||
success = await forward_to_external(endpoint, payload, api_key if api_key else None, http_method)
|
||||
|
||||
duration = asyncio.get_event_loop().time() - start_time
|
||||
message_duration.labels(endpoint=endpoint).observe(duration)
|
||||
|
||||
if success:
|
||||
# Acknowledge message (remove from queue)
|
||||
await msg.ack()
|
||||
messages_processed.labels(status="success", endpoint=endpoint).inc()
|
||||
print(f"✅ Message processed and forwarded successfully for {endpoint}")
|
||||
else:
|
||||
# Negative acknowledge (requeue for retry)
|
||||
await msg.nak()
|
||||
messages_processed.labels(status="failed", endpoint=endpoint).inc()
|
||||
print(f"❌ Failed to forward message for {endpoint}, requeued for retry")
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"❌ Invalid JSON in message: {e}")
|
||||
await msg.term() # Terminate (don't retry)
|
||||
messages_processed.labels(status="invalid", endpoint=endpoint).inc()
|
||||
except Exception as e:
|
||||
print(f"❌ Error processing message for {endpoint}: {e}")
|
||||
await msg.nak() # Requeue for retry
|
||||
messages_processed.labels(status="error", endpoint=endpoint).inc()
|
||||
finally:
|
||||
messages_in_flight.dec()
|
||||
|
||||
async def consume_messages():
|
||||
"""Consume messages from NATS JetStream"""
|
||||
global nc, js
|
||||
|
||||
try:
|
||||
print(f"Connecting to NATS at {NATS_URL}...")
|
||||
nc = await nats.connect(
|
||||
servers=[NATS_URL],
|
||||
user=NATS_USER,
|
||||
password=NATS_PASSWORD,
|
||||
reconnect_time_wait=2,
|
||||
max_reconnect_attempts=-1 # Retry forever
|
||||
)
|
||||
js = nc.jetstream()
|
||||
print("✅ Connected to NATS JetStream")
|
||||
|
||||
# Subscribe to messages (using wildcard to catch all API subjects)
|
||||
print(f"Subscribing to stream '{NATS_STREAM}', consumer '{NATS_CONSUMER}'...")
|
||||
print(f" Subject pattern: {NATS_SUBJECT}")
|
||||
print(f" Endpoints configured: {len(ENDPOINT_MAPPING)}")
|
||||
|
||||
psub = await js.pull_subscribe(
|
||||
NATS_SUBJECT,
|
||||
NATS_CONSUMER,
|
||||
stream=NATS_STREAM
|
||||
)
|
||||
print(f"✅ Subscribed, waiting for messages...")
|
||||
|
||||
# Start metrics server
|
||||
start_http_server(9090)
|
||||
print("✅ Metrics server started on port 9090")
|
||||
|
||||
# Process messages in a loop
|
||||
while running:
|
||||
try:
|
||||
# Fetch messages (batch of up to 10)
|
||||
msgs = await psub.fetch(10, timeout=5)
|
||||
for msg in msgs:
|
||||
# Process asynchronously
|
||||
asyncio.create_task(process_message(msg))
|
||||
except asyncio.TimeoutError:
|
||||
# No messages available, continue
|
||||
continue
|
||||
except Exception as e:
|
||||
print(f"⚠️ Error fetching messages: {e}")
|
||||
await asyncio.sleep(1)
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Fatal error: {e}")
|
||||
raise
|
||||
finally:
|
||||
if nc:
|
||||
await nc.close()
|
||||
print("NATS connection closed")
|
||||
|
||||
async def main():
|
||||
"""Main entry point"""
|
||||
print("🚀 Starting NATS Worker (Multi-Endpoint)...")
|
||||
print(f" External Base URL: {BASE_URL}")
|
||||
print(f" Concurrency: {WORKER_CONCURRENCY}")
|
||||
print(f" Retry Attempts: {RETRY_ATTEMPTS}")
|
||||
print(f" Configured Endpoints:")
|
||||
for fastapi_endpoint, external_url in ENDPOINT_MAPPING.items():
|
||||
print(f" {fastapi_endpoint} → {external_url}")
|
||||
|
||||
try:
|
||||
await consume_messages()
|
||||
except KeyboardInterrupt:
|
||||
print("\n🛑 Worker stopped by user")
|
||||
except Exception as e:
|
||||
print(f"❌ Worker crashed: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
||||
Reference in New Issue
Block a user