main.py

"""
main.py
=======
Production-style entry point demonstrating all system components together.

Run: python main.py
"""

from __future__ import annotations

import logging
import random
import time

from core.task       import Task, TaskPriority
from core.queue      import PriorityTaskQueue, DeadLetterQueue, QueueFactory
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, TagAuthMiddleware, build_pipeline
from plugins.storage     import InMemoryStorage

# ---------------------------------------------------------------------------
# Logging configuration (structured, like production)
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("main")


# ---------------------------------------------------------------------------
# Simulated workloads
# ---------------------------------------------------------------------------
def payment_processor(order_id: str, amount: float) -> dict:
    """Simulate a payment API call with occasional transient failures."""
    time.sleep(random.uniform(0.05, 0.2))
    if random.random() < 0.2:          # 20% transient failure rate
        raise ConnectionError(f"Payment gateway timeout for order {order_id}")
    return {"order_id": order_id, "amount": amount, "status": "charged"}


def email_sender(recipient: str, subject: str) -> dict:
    """Simulate sending an email — fast, rarely fails."""
    time.sleep(random.uniform(0.01, 0.05))
    if random.random() < 0.05:
        raise RuntimeError(f"SMTP relay error for {recipient}")
    return {"to": recipient, "subject": subject, "delivered": True}


def report_generator(report_id: str) -> dict:
    """Simulate a CPU-bound report — slow, reliable."""
    time.sleep(random.uniform(0.1, 0.4))
    return {"report_id": report_id, "rows": random.randint(100, 10_000)}


# ---------------------------------------------------------------------------
def main() -> None:
    print("\n" + "═" * 60)
    print("  Distributed Task Queue System  —  FAANG-Level Demo")
    print("═" * 60 + "\n")

    # ── Infrastructure ──────────────────────────────────────────────
    metrics    = MetricsCollector()
    storage    = InMemoryStorage()
    queue      = PriorityTaskQueue(maxsize=500)
    dlq        = DeadLetterQueue()

    def persist_on_success(task):
        storage.save(task)
        metrics.on_success(task)

    def persist_on_failure(task):
        storage.save(task)
        metrics.on_failure(task)

    def persist_on_retry(task):
        storage.save(task)
        metrics.on_retry(task)

    pool = WorkerPool(
        num_workers=6,
        on_success=persist_on_success,
        on_failure=persist_on_failure,
        on_retry=persist_on_retry,
    )

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

    middleware = build_pipeline(
        RateLimiter(rate=50.0),            # token bucket: 50 tasks/sec
        TagAuthMiddleware({"blocked"}),    # reject tasks tagged "blocked"
        StructuredLogger(),                # emit structured log lines
    )

    dispatcher = Dispatcher(
        queue=queue,
        pool=pool,
        scheduler=scheduler,
        middleware=middleware,
        dlq=dlq,
        poll_interval=0.01,
        max_inflight=20,
    )

    # ── Enqueue workload ────────────────────────────────────────────
    logger.info("Enqueuing 60 tasks across 3 priority levels …")

    # CRITICAL: payment tasks
    for i in range(15):
        queue.enqueue(Task(
            name=f"payment_{i:03d}",
            func=payment_processor,
            args=(f"ORD-{1000+i}", round(random.uniform(9.99, 499.99), 2)),
            priority=TaskPriority.CRITICAL,
            max_retries=4,
            tags=["payments"],
        ))

    # HIGH: email tasks
    for i in range(25):
        queue.enqueue(Task(
            name=f"email_{i:03d}",
            func=email_sender,
            args=(f"user_{i}@example.com", f"Your order #{1000+i} is confirmed"),
            priority=TaskPriority.HIGH,
            max_retries=3,
            tags=["notifications"],
        ))

    # LOW: report tasks
    for i in range(20):
        queue.enqueue(Task(
            name=f"report_{i:03d}",
            func=report_generator,
            args=(f"RPT-{i:04d}",),
            priority=TaskPriority.LOW,
            max_retries=2,
            tags=["reports"],
        ))

    logger.info(f"Queue depth: {queue.size()} tasks")

    # ── Dispatch ────────────────────────────────────────────────────
    start = time.monotonic()

    with dispatcher:
        # Poll until queue is drained or timeout
        deadline = start + 30.0
        while not queue.is_empty() and time.monotonic() < deadline:
            time.sleep(0.1)
        time.sleep(1.0)    # allow in-flight tasks to settle

    pool.shutdown(wait=True)
    elapsed = time.monotonic() - start

    # ── Report ──────────────────────────────────────────────────────
    report = metrics.report()

    print("\n" + "─" * 60)
    print("  FINAL METRICS REPORT")
    print("─" * 60)

    t = report["totals"]
    print(f"  Tasks submitted : {t['total']}")
    print(f"  ✓ Succeeded     : {t['success']}")
    print(f"  ✗ Failed        : {t['failure']}")
    print(f"  ↻ Retried       : {t['retried']}")
    print(f"  ☠ Dead-lettered : {t['dead']}")
    print(f"  Success rate    : {report['success_rate_%']}%")

    lat = report["latency_ms"]
    print(f"\n  Latency (last 60s window)")
    print(f"    p50  : {lat['p50']} ms")
    print(f"    p95  : {lat['p95']} ms")
    print(f"    p99  : {lat['p99']} ms")

    tp = report["throughput"]
    print(f"\n  Throughput      : {tp['tasks_per_sec']} tasks/sec")
    print(f"  Wall-clock time : {elapsed:.2f}s")

    # Storage stats
    print(f"\n  Persisted tasks : {storage.count()}")
    status_breakdown = storage.stats_by_status()
    for status, count in sorted(status_breakdown.items()):
        print(f"    {status:<10}: {count}")

    # DLQ report
    dlq_summary = dlq.summary()
    print(f"\n  Dead-Letter Queue")
    print(f"    Total dead     : {dlq_summary['total_dead']}")
    if dlq_summary["error_breakdown"]:
        print(f"    Error types    :")
        for err, cnt in dlq_summary["error_breakdown"].items():
            print(f"      [{cnt}x] {err[:60]}")

    print("\n" + "═" * 60 + "\n")


if __name__ == "__main__":
    main()