Script to benchmark and compare performance of Lua script in Redis

redis_benchmark_script.py

import redis
import time
import random
import string
import json
from typing import Dict
import matplotlib.pyplot as plt

def generate_random_string(length: int) -> str:
    """Generate a random string of fixed length"""
    return ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))

class URLShortenerBenchmark:
    def __init__(self, host='localhost', port=6380, db=0):
        self.redis = redis.Redis(host=host, port=port, db=db, decode_responses=True)
        self.redis.flushdb()
        
        self.get_user_urls_script = self.redis.register_script("""
            local short_urls = redis.call('SMEMBERS', KEYS[1])
            local result = {}
            for i, short_url in ipairs(short_urls) do
                local url_key = 'url:short:' .. short_url
                local long_url = redis.call('HGET', url_key, 'long_url')
                result[i] = {short_url, long_url}
            end
            return cjson.encode(result)
        """)

    def setup_test_data(self, email: str, num_urls: int):
        """Setup test data with specified number of URLs"""
        for _ in range(num_urls):
            short_url = generate_random_string(6)
            long_url = f"https://example.com/{generate_random_string(10)}"
            
            self.redis.hset(
                f"url:short:{short_url}",
                mapping={"long_url": long_url}
            )
            self.redis.sadd(f"url:user:{email}", short_url)

    def get_urls_n_plus_one(self, email: str) -> Dict[str, str]:
        """Traditional N+1 approach"""
        result = {}
        short_urls = self.redis.smembers(f"url:user:{email}")
        
        for short_url in short_urls:
            long_url = self.redis.hget(f"url:short:{short_url}", "long_url")
            result[short_url] = long_url
            
        return result

    def get_urls_lua(self, email: str) -> Dict[str, str]:
        """Lua script approach"""
        result = self.get_user_urls_script(keys=[f"url:user:{email}"])
        url_pairs = json.loads(result)
        return {pair[0]: pair[1] for pair in url_pairs}

    def run_benchmark(self, url_counts: list, iterations: int = 5) -> Dict:
        """Run benchmark for different URL counts"""
        results = {
            'n_plus_one': [],
            'lua_script': []
        }
        
        email = "test@example.com"
        
        for num_urls in url_counts:
            print(f"\nBenchmarking {num_urls} URLs...")
            
            self.redis.flushdb()
            
            print(f"Setting up test data...")
            self.setup_test_data(email, num_urls)
            
            print(f"Running N+1 queries benchmark...")
            n_plus_one_times = []
            for i in range(iterations):
                start_time = time.time()
                self.get_urls_n_plus_one(email)
                n_plus_one_times.append(time.time() - start_time)
                print(f"  Iteration {i+1}/{iterations}", end='\r')
            print()
            

            print(f"Running Lua script benchmark...")
            lua_times = []
            for i in range(iterations):
                start_time = time.time()
                self.get_urls_lua(email)
                lua_times.append(time.time() - start_time)
                print(f"  Iteration {i+1}/{iterations}", end='\r')
            print()
            
            results['n_plus_one'].append(sum(n_plus_one_times) / iterations)
            results['lua_script'].append(sum(lua_times) / iterations)
        
        return results

def plot_results(url_counts: list, results: Dict):
    """Plot benchmark results"""
    plt.figure(figsize=(10, 6))
    plt.plot(url_counts, results['n_plus_one'], 'o-', label='N+1 Queries')
    plt.plot(url_counts, results['lua_script'], 'o-', label='Lua Script')
    plt.xlabel('Number of URLs')
    plt.ylabel('Time (seconds)')
    plt.title('Redis Query Performance: N+1 vs Lua Script')
    plt.legend()
    plt.grid(True)
    plt.show()

if __name__ == "__main__":
    benchmark = URLShortenerBenchmark()
    url_counts = [10, 50, 100, 500, 1000, 5000]
    results = benchmark.run_benchmark(url_counts)
    
    print("\nBenchmark Results (seconds):")
    print("\nURLs Count | N+1 Queries | Lua Script")
    print("-" * 40)
    for i, count in enumerate(url_counts):
        print(f"{count:>10} | {results['n_plus_one'][i]:>11.4f} | {results['lua_script'][i]:>9.4f}")
    
    plot_results(url_counts, results)