worker_daemon.py

"""
worker_daemon.py
================
Standalone worker process — no HTTP server, just task processing.

In a microservices deployment you run multiple replicas of this daemon
alongside a single API server. This separation means you can scale
workers independently of the API layer (a common FAANG interview design point).

Architecture
------------
   [Client] ──HTTP──► [API container]  ──ZADD──► [Redis]
                                                    │
                                                  ZPOPMIN
                                                    │
                                                    ▼
                                          [Worker container ×N]
                                          (this file, replicated)

Run locally
-----------
    python worker_daemon.py

In Docker
---------
    docker-compose up worker
"""

from __future__ import annotations

import logging
import os
import random
import signal
import sys
import time

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] worker — %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("worker_daemon")

# ── Job registry (same as api/main.py) ───────────────────────────────────

def job_payment(order_id: str, amount: float) -> dict:
    time.sleep(random.uniform(0.05, 0.2))
    if random.random() < 0.15:
        raise ConnectionError(f"Payment gateway timeout for order {order_id}")
    return {"order_id": order_id, "amount": amount, "charged": True}

def job_email(recipient: str, subject: str, body: str = "") -> dict:
    time.sleep(random.uniform(0.01, 0.06))
    if random.random() < 0.05:
        raise RuntimeError(f"SMTP relay error for {recipient}")
    return {"to": recipient, "subject": subject, "delivered": True}

def job_report(report_id: str, rows: int = 1000) -> dict:
    time.sleep(random.uniform(0.1, 0.4))
    return {"report_id": report_id, "rows": rows, "status": "generated"}

def job_sync(resource: str) -> dict:
    time.sleep(random.uniform(0.05, 0.15))
    return {"resource": resource, "synced": True}

JOB_REGISTRY = {
    "payment": job_payment,
    "email":   job_email,
    "report":  job_report,
    "sync":    job_sync,
}


def main() -> None:
    redis_url   = os.getenv("REDIS_URL", "redis://localhost:6379/0")
    num_workers = int(os.getenv("WORKERS", "4"))

    logger.info("═" * 50)
    logger.info("  Task Queue Worker Daemon starting")
    logger.info(f"  Redis  : {redis_url}")
    logger.info(f"  Workers: {num_workers}")
    logger.info("═" * 50)

    from core.queue      import DeadLetterQueue
    from core.worker     import WorkerPool
    from core.scheduler  import Scheduler, ExponentialBackoffWithJitter
    from core.dispatcher import Dispatcher
    from monitoring.metrics import MetricsCollector
    from plugins.middleware  import RateLimiter, StructuredLogger, build_pipeline

    try:
        from plugins.redis_backend import RedisQueue, RedisStorage, RedisDLQ
        queue   = RedisQueue(redis_url=redis_url, job_registry=JOB_REGISTRY)
        storage = RedisStorage(redis_url=redis_url)
        dlq     = RedisDLQ(redis_url=redis_url)
        logger.info("✓ Connected to Redis")
    except Exception as e:
        logger.warning(f"Redis unavailable ({e}), using in-memory fallback")
        from core.queue import PriorityTaskQueue
        from plugins.storage import InMemoryStorage
        queue   = PriorityTaskQueue(maxsize=1000)
        storage = InMemoryStorage()
        dlq     = DeadLetterQueue()

    mc = MetricsCollector()

    pool = WorkerPool(
        num_workers = num_workers,
        on_success  = lambda t: (storage.save(t), mc.on_success(t)),
        on_failure  = lambda t: (storage.save(t), mc.on_failure(t)),
        on_retry    = lambda t: (storage.save(t), mc.on_retry(t)),
    )

    scheduler = Scheduler(
        enqueue_fn = queue.enqueue,
        strategy   = ExponentialBackoffWithJitter(base=0.5, cap=30.0),
        dlq_fn     = dlq.enqueue,
    )

    dispatcher = Dispatcher(
        queue         = queue,
        pool          = pool,
        scheduler     = scheduler,
        middleware    = build_pipeline(RateLimiter(200.0), StructuredLogger()),
        dlq           = dlq,
        poll_interval = 0.05,
        max_inflight  = num_workers * 2,
    )

    # ── Graceful shutdown on SIGTERM / SIGINT ─────────────────
    def shutdown(signum, frame):
        logger.info(f"Signal {signum} received — shutting down gracefully …")
        dispatcher.stop(timeout=10.0)
        pool.shutdown(wait=True)
        report = mc.report()
        logger.info(f"Final stats: {report['totals']}")
        sys.exit(0)

    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT,  shutdown)

    # ── Run forever ───────────────────────────────────────────
    with dispatcher:
        logger.info("✓ Dispatcher running — waiting for tasks …")
        while True:
            time.sleep(10)
            report = mc.report()
            logger.info(
                f"[heartbeat] done={report['totals']['success']} "
                f"failed={report['totals']['failure']} "
                f"retried={report['totals']['retried']} "
                f"queue_depth={queue.size()} "
                f"tps={report['throughput']['tasks_per_sec']}"
            )


if __name__ == "__main__":
    main()