redis_postgres_benchmark/main.py

#!/usr/bin/env python3
"""
Redis vs PostgreSQL Read Benchmark - Fair Comparison
"""
import os
import sys
import time
import random
import logging
import csv
import redis
import psycopg2
from locust import User, task, events

# ─── Logging ───
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(message)s',
    handlers=[
        logging.FileHandler('debug.log'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger()

# ─── Config ───
TOTAL_RECORDS = 125_000_000
PHONE_MIN = 100_000_000
PHONE_MAX = 224_999_999

REDIS_CONFIG = {
    "host": "localhost",
    "port": 6379,
    "db": 0,
    "decode_responses": False,
    "socket_connect_timeout": 5,
    "socket_timeout": 10,
}

POSTGRES_CONFIG = {
    "host": "localhost",
    "port": 5432,
    "database": "numbers_db",
    "user": "myuser",
    "password": os.environ.get("POSTGRES_PASSWORD", "mypassword"),
}


# ─────────────────────────────────────────────
# Redis Benchmark - با GET
# ─────────────────────────────────────────────
class RedisReadUser(User):
    _pool = None

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if RedisReadUser._pool is None:
            RedisReadUser._pool = redis.ConnectionPool(**REDIS_CONFIG)
        # هر User یک client اختصاصی از pool می‌گیره
        self.client = redis.Redis(connection_pool=RedisReadUser._pool)

    @task
    def get_random(self):
        phone_number = random.randint(PHONE_MIN, PHONE_MAX)
        start = time.perf_counter()
        try:
            four_digit_code = self.client.get(phone_number)
            elapsed = (time.perf_counter() - start) * 1000
            if four_digit_code:
                print(f"Redis | Phone: {phone_number} | Code: {four_digit_code} | Time: {elapsed:.3f}ms")
            else:
                print(f"Redis | Phone: {phone_number} | Code: NOT FOUND | Time: {elapsed:.3f}ms")
            self.environment.events.request.fire(
                request_type="GET",
                name="redis_get",
                response_time=elapsed,
                response_length=len(four_digit_code) if four_digit_code else 0,
                exception=None,
                context={},
            )
        except Exception as e:
            elapsed = (time.perf_counter() - start) * 1000
            print(f"❌ Redis ERROR: {e} | Phone: {phone_number} | Time: {elapsed:.3f}ms")
            self.environment.events.request.fire(
                request_type="GET",
                name="redis_get",
                response_time=elapsed,
                response_length=0,
                exception=e,
                context={},
            )


# ─────────────────────────────────────────────
# PostgreSQL Benchmark - کانکشن اختصاصی برای هر User
# ─────────────────────────────────────────────
class PostgresReadUser(User):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # ✅ هر User instance یک کانکشن اختصاصی دارد (مثل Redis)
        self._connection = psycopg2.connect(
            host=POSTGRES_CONFIG['host'],
            port=POSTGRES_CONFIG['port'],
            database=POSTGRES_CONFIG['database'],
            user=POSTGRES_CONFIG['user'],
            password=POSTGRES_CONFIG['password']
        )
        # تنظیمات برای عملکرد بهتر
        self._connection.set_session(autocommit=True)

    def on_stop(self):
        """آزادسازی کانکشن هنگام توقف User"""
        if self._connection is not None and not self._connection.closed:
            self._connection.close()
            self._connection = None

    @task
    def get_random(self):
        phone_number = random.randint(PHONE_MIN, PHONE_MAX)
        start = time.perf_counter()
        try:
            with self._connection.cursor() as cur:
                cur.execute(
                    "SELECT random_code FROM phone_numbers WHERE phone_number = %s",
                    (phone_number,)
                )
                result = cur.fetchone()
            elapsed = (time.perf_counter() - start) * 1000

            if result:
                random_code = result[0]
                print(f"Postgres | Phone: {phone_number} | Code: {random_code} | Time: {elapsed:.3f}ms")
            else:
                print(f"Postgres | Phone: {phone_number} | Code: NOT FOUND | Time: {elapsed:.3f}ms")

            self.environment.events.request.fire(
                request_type="GET",
                name="postgres_select",
                response_time=elapsed,
                response_length=4,
                exception=None,
                context={},
            )
        except Exception as e:
            elapsed = (time.perf_counter() - start) * 1000
            print(f"❌ Postgres ERROR: {e} | Phone: {phone_number} | Time: {elapsed:.3f}ms")
            self.environment.events.request.fire(
                request_type="GET",
                name="postgres_select",
                response_time=elapsed,
                response_length=0,
                exception=e,
                context={},
            )


# ─────────────────────────────────────────────
# اجرا
# ─────────────────────────────────────────────
if __name__ == "__main__":
    print(f"""
╔══════════════════════════════════════════════════════════════╗
║        Redis vs PostgreSQL READ Benchmark (125M)             ║
╠══════════════════════════════════════════════════════════════╣
║                                                              ║
║  Total Records: {TOTAL_RECORDS:,}                            ║
║                                                              ║
║                                                              ║
║  ⚠️  Response Time = فقط زمان Query (بدون Connection)        ║
╚══════════════════════════════════════════════════════════════╝
    """)