redis_postgres_benchmark/load_redis.py

import os
import redis
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, as_completed
import time
import random

# ─────────────────────────────────────────────
# کانفیگ Redis
# ─────────────────────────────────────────────
REDIS_CONFIG = {
    "host": "localhost",
    "port": 6379,
    "db": 0,
    "decode_responses": True,
    "socket_connect_timeout": 5,
    "socket_keepalive": True,
}

NUM_WORKERS = 8
BATCH_SIZE = 100_000

# ═══════════════════════════════════════════════════════════════
# بازه اعداد - مطابق با داده‌های لود شده
# ═══════════════════════════════════════════════════════════════
TOTAL_RECORDS = 125_000_000
PHONE_START = 100_000_000
PHONE_END = PHONE_START + TOTAL_RECORDS - 1  # 224,999,999


def optimize_redis(r):
    """تنظیمات بهینه Redis"""
    r.config_set("maxmemory", "16gb")
    r.config_set("maxmemory-policy", "allkeys-lru")
    r.config_set("save", "")
    print("✓ Redis optimized")


def load_chunk_to_redis(args):
    """لود یک بخش از بازه اعداد به Redis"""
    worker_id, start_phone, end_phone = args
    r = redis.Redis(**REDIS_CONFIG)
    try:
        print(f"[Worker {worker_id}] Starting... phone: {start_phone:,} to {end_phone:,}")

        loaded = 0
        batch_count = 0
        start_time = time.time()

        pipe = r.pipeline()
        current_phone = start_phone

        while current_phone <= end_phone:
            random_value = random.randint(1000, 9999)
            pipe.set(str(current_phone), random_value)
            loaded += 1
            current_phone += 1

            if loaded >= BATCH_SIZE * (batch_count + 1):
                pipe.execute()
                batch_count += 1
                elapsed = time.time() - start_time
                rate = loaded / elapsed if elapsed > 0 else 0
                print(f"[Worker {worker_id}] Batch {batch_count} done ({loaded:,} rows, {rate:,.0f} rows/sec)")

        if len(pipe.command_stack) > 0:
            pipe.execute()

        elapsed = time.time() - start_time
        rate = loaded / elapsed if elapsed > 0 else 0
        print(f"[Worker {worker_id}] Done! {loaded:,} rows in {elapsed:.1f}s ({rate:,.0f} rows/sec)")
        return loaded

    except Exception as e:
        print(f"[Worker {worker_id}] ERROR: {e}")
        raise
    finally:
        r.close()


def main():
    print("=" * 60)
    print("Redis Bulk Loader - Sequential Phone Numbers")
    print("=" * 60)
    print(f"Phone Range: {PHONE_START:,} to {PHONE_END:,}")
    print(f"Total Records: {TOTAL_RECORDS:,}\n")

    num_workers = min(NUM_WORKERS, mp.cpu_count())
    print(f"Using {num_workers} workers\n")

    r = redis.Redis(**REDIS_CONFIG)
    try:
        optimize_redis(r)
        print("Clearing existing keys...")
        r.flushdb()
    finally:
        r.close()

    start_time = time.time()

    try:
        # ═══════════════════════════════════════════════════════
        # تقسیم بازه اعداد بین workers
        # ═══════════════════════════════════════════════════════
        records_per_worker = TOTAL_RECORDS // num_workers
        args = []

        for i in range(num_workers):
            if i < num_workers - 1:
                start_phone = PHONE_START + (i * records_per_worker)
                end_phone = PHONE_START + ((i + 1) * records_per_worker) - 1
            else:
                # آخرین worker باقیمانده را هم می‌گیرد
                start_phone = PHONE_START + (i * records_per_worker)
                end_phone = PHONE_END

            args.append((i, start_phone, end_phone))
            print(f"Worker {i}: {start_phone:,} to {end_phone:,} ({end_phone - start_phone + 1:,} records)")

        print("\nStarting workers...\n")

        with ProcessPoolExecutor(max_workers=num_workers) as executor:
            futures = [executor.submit(load_chunk_to_redis, arg) for arg in args]
            total = 0

            for future in as_completed(futures):
                try:
                    total += future.result()
                except Exception as e:
                    print(f"Worker failed: {e}")

                elapsed = time.time() - start_time
                rate = total / elapsed if elapsed > 0 else 0
                remaining = (TOTAL_RECORDS - total) / rate if rate > 0 else 0
                percent = (total * 100) // TOTAL_RECORDS

                print(f"\n>>> Progress: {total:,}/{TOTAL_RECORDS:,} ({percent}%) | "
                      f"Speed: {rate:,.0f} rows/sec | "
                      f"ETA: {remaining / 60:.1f} min")

        print(f"\n✓ Total rows loaded: {total:,}")

        # بررسی نهایی
        r = redis.Redis(**REDIS_CONFIG)
        try:
            print("\n--- Sample Results ---")
            sample_phones = [
                PHONE_START,
                PHONE_START + 10_000_000,
                PHONE_START + 50_000_000,
                PHONE_END
            ]
            for phone in sample_phones:
                value = r.get(str(phone))
                print(f"  {phone:,} → {value}")

            final_count = r.dbsize()
            print(f"\n✓ Final count in Redis: {final_count:,}")
        finally:
            r.close()

        total_time = time.time() - start_time
        print(f"\n✓ Done! Total time: {total_time / 60:.1f} minutes")
        print(f"✓ Average speed: {TOTAL_RECORDS / total_time:,.0f} rows/sec")

    except Exception as e:
        print(f"\n✗ Error: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()