kubernetes/manifests/core/worker-script.yaml

apiVersion: v1
kind: ConfigMap
metadata:
  name: worker-script
  namespace: core
data:
  worker.py: |
    #!/usr/bin/env python3
    """
    NATS Worker - "New Idea" Implementation
    - Single Domain per Worker (Stream + Consumer)
    - Persistent HTTP Session (Connection Pooling)
    - JetStream Native Retries (NAK)
    - No internal retry loops
    """
    import os
    import json
    import asyncio
    import aiohttp
    import nats
    from nats.errors import TimeoutError as NatsTimeoutError
    from nats.js.api import ConsumerConfig
    from prometheus_client import Counter, Histogram, Gauge, start_http_server
    import signal
    import sys
    from typing import Any, Dict, List, Union, Tuple

    # 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')

    # Configuration
    NATS_URL = os.getenv("NATS_URL", "nats://nats-server:4222")
    NATS_USER = os.getenv("NATS_USER", "admin")
    NATS_PASSWORD = os.getenv("NATS_PASSWORD", "")

    # Domain Configuration (Env Vars from Deployment)
    NATS_STREAM = os.getenv("NATS_STREAM", "ORDERS")
    NATS_CONSUMER = os.getenv("NATS_CONSUMER", "orders-worker")
    # Supports comma-separated subjects: "api.v1.orders.*,api.v2.orders.*"
    FILTER_SUBJECTS = os.getenv("FILTER_SUBJECT", "api.v1.mob.orders.*").split(",")

    WORKER_CONCURRENCY = int(os.getenv("WORKER_CONCURRENCY", "10"))
    BASE_URL = os.getenv("EXTERNAL_BASE_URL", "https://jupiter.nearle.app")
    # Must stay comfortably above the outbound HTTP timeout below (30s), otherwise
    # JetStream can redeliver a message to another worker while this one is still
    # waiting on a slow-but-legitimate response, causing a duplicate forward.
    ACK_WAIT_SECONDS = int(os.getenv("ACK_WAIT_SECONDS", "60"))

    # Endpoint mapping is still useful for constructing the target URL
    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",
        "/live/api/v1/mob/orders/createorder": f"{BASE_URL}/live/api/v1/orders/createorder",
        "/live/api/v1/web/products/create": f"{BASE_URL}/live/api/v1/products/create",
        "/live/api/v1/mob/customers/login": f"{BASE_URL}/live/api/v1/customers/login",
        "/live/api/v1/mob/customers/create": f"{BASE_URL}/live/api/v1/customers/create",
    }

    # Global State
    nc = None
    js = None
    session: aiohttp.ClientSession = None
    running = True
    semaphore = None # Initialized in main
    active_tasks: set = set() # In-flight process_message() tasks, drained on shutdown

    def signal_handler(sig, frame):
        global running
        print("\n?? Shutdown signal received...")
        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") -> tuple[str, Any]:
        """
        Forward to external API using global session. 
        Returns (status_code, response_data).
        Does NOT retry internally.
        """
        external_url = ENDPOINT_MAPPING.get(endpoint)
        if not external_url:
            print(f"?? Unknown endpoint: {endpoint}, dropping.")
            return "DROP", None

        headers = {"Content-Type": "application/json"}
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"

        if "data" not in payload:
            print(f"?? Message for {endpoint} missing 'data' field, dropping.")
            return "DROP", None
        data_to_forward = payload["data"]

        # Payload Normalization for logs
        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):
                data_to_forward = [data_to_forward]

        http_method = method.upper() if method else "POST"
    
        try:
            # We use the global 'session' here
            async with session.request(
                method=http_method, 
                url=external_url, 
                json=data_to_forward if http_method != "GET" else None,
                params=data_to_forward if http_method == "GET" else None,
                headers=headers, 
                timeout=aiohttp.ClientTimeout(total=30)
            ) as response:
            
                body = None
                try:
                    body = await response.json()
                except:
                    body = await response.text()

                if 200 <= response.status < 300:
                    return "OK", body
                elif response.status >= 500:
                    print(f"?? Server Error {response.status} from {endpoint}")
                    return "RETRY", body
                else:
                    print(f"?? Client Error {response.status} from {endpoint}: {body}")
                    return "DROP", body # 4xx errors should generally not be retried infinitely

        except asyncio.TimeoutError:
            print(f"?? Timeout calling {endpoint}")
            return "RETRY", None
        except Exception as e:
            print(f"?? Network/Client Error calling {endpoint}: {e}")
            return "RETRY", None

    async def process_message(msg):
        async with semaphore:
            messages_in_flight.inc()
            start_time = asyncio.get_event_loop().time()
            endpoint = "unknown"
        
            try:
                payload = json.loads(msg.data.decode())
                endpoint = payload.get("endpoint", payload.get("original_path", "unknown"))
                http_method = payload.get("method", "POST")

                print(f"?? Processing {endpoint}")

                api_key = os.getenv("EXTERNAL_ENDPOINT_API_KEY", "")
            
                # Single attempt
                status, response_data = await forward_to_external(endpoint, payload, api_key, http_method)
            
                if status == "OK":
                    await msg.ack()
                    messages_processed.labels(status="success", endpoint=endpoint).inc()
                
                    # Request-Reply Logic
                    if msg.reply:
                        reply_payload = json.dumps(response_data) if isinstance(response_data, (dict, list)) else str(response_data)
                        await nc.publish(msg.reply, reply_payload.encode())

                elif status == "DROP":
                    # Unrecoverable error or unknown endpoint
                    await msg.term() # Terminate stops redelivery
                    messages_processed.labels(status="dropped", endpoint=endpoint).inc()
            
                else: # RETRY
                    # Let JetStream handle backoff
                    await msg.nak(delay=2) # Custom delay before redelivery if desired, or just nak()
                    messages_processed.labels(status="retried", endpoint=endpoint).inc()

            except json.JSONDecodeError:
                print("?? Invalid JSON, terminating message")
                await msg.term()
            except Exception as e:
                print(f"? Critical Worker Error: {e}")
                # Delay avoids a tight redelivery loop if this is a persistent bug
                await msg.nak(delay=5)
            finally:
                messages_in_flight.dec()
                duration = asyncio.get_event_loop().time() - start_time
                message_duration.labels(endpoint=endpoint).observe(duration)

    async def main():
        global nc, js, session, semaphore
    
        print(f"?? Starting Worker for Domain: {NATS_CONSUMER}")
        print(f"?? Stream: {NATS_STREAM}, Subjects: {FILTER_SUBJECTS}")
        print(f"?? Concurrency: {WORKER_CONCURRENCY}")

        semaphore = asyncio.Semaphore(WORKER_CONCURRENCY)
        session = aiohttp.ClientSession()

        try:
            nc = await nats.connect(
                servers=[NATS_URL],
                user=NATS_USER,
                password=NATS_PASSWORD,
                reconnect_time_wait=2,
                max_reconnect_attempts=-1
            )
            js = nc.jetstream()
            print("? Connected to NATS JetStream")

            # Create Pull Subscriptions for each filter subject
            # All sharing the same Consumer Name ensuring load balancing if multiple pods run this
        
            subs = []
            for subject in FILTER_SUBJECTS:
                subject = subject.strip()
                if not subject: continue
            
                # FIX: On WorkQueue streams, we cannot reuse the same durable name for different filters
                # We append a sanitized version of the subject to ensure uniqueness per filter
                clean_subject_suffix = subject.replace(".", "_").replace("*", "all").replace(">", "all")
                durable_name = f"{NATS_CONSUMER}_{clean_subject_suffix}"
            
                print(f"?? Subscribing to {subject} on durable consumer '{durable_name}'")
                try:
                    sub = await js.pull_subscribe(
                        subject,
                        durable=durable_name,
                        stream=NATS_STREAM,
                        config=ConsumerConfig(ack_wait=ACK_WAIT_SECONDS)
                    )
                    subs.append(sub)
                except Exception as e:
                    print(f"? Failed to subscribe to {subject}: {e}")

            if not subs:
                print("? No active subscriptions!")
                return

            start_http_server(9090)
            print("?? Metrics on :9090")

            while running:
                # Poll all subscriptions
                for sub in subs:
                    try:
                        # Batch size 10, short timeout to keep loop responsive
                        msgs = await sub.fetch(10, timeout=0.5)
                        for m in msgs:
                            task = asyncio.create_task(process_message(m))
                            active_tasks.add(task)
                            task.add_done_callback(active_tasks.discard)
                    except NatsTimeoutError:
                        pass
                    except Exception as e:
                        print(f"?? Fetch Error: {e}")
                        await asyncio.sleep(1)
            
                # Small sleep to prevent tight loop if no messages
                # await asyncio.sleep(0.01) 

        except Exception as e:
            print(f"?? Fatal Error: {e}")
        finally:
            if active_tasks:
                print(f"?? Draining {len(active_tasks)} in-flight message(s) before shutdown...")
                await asyncio.wait(active_tasks, timeout=25)
            if session:
                await session.close()
            if nc:
                await nc.close()
            print("?? Worker Shutdown")

    if __name__ == "__main__":
        asyncio.run(main())