load_generator/generator.py

import os
import asyncio
import time
import httpx
import redis.asyncio as redis

TARGET_URL = os.getenv("TARGET_URL", "http://service-a:8000/request")
REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")

# Keep a set of strong references to running tasks so they aren't garbage collected
running_tasks = set()

async def send_single_request(client):
    try:
        await client.get(TARGET_URL)
    except Exception:
        pass

async def main():
    print(f"Starting load generator targeting: {TARGET_URL}")
    
    # Wait for service-a to be ready before firing
    await asyncio.sleep(5.0)
    
    redis_client = redis.from_url(REDIS_URL, decode_responses=True)
    await redis_client.set("load_generator:target_rps", "50") # Default starting RPS
    
    async with httpx.AsyncClient(timeout=3.0) as client:
        while True:
            try:
                rps_str = await redis_client.get("load_generator:target_rps")
                target_rps = int(rps_str) if rps_str else 50
            except Exception:
                target_rps = 50
                
            if target_rps <= 0:
                await asyncio.sleep(1.0)
                continue
                
            start_time = time.perf_counter()
            interval = 1.0 / target_rps
            
            for _ in range(target_rps):
                task = asyncio.create_task(send_single_request(client))
                running_tasks.add(task)
                task.add_done_callback(running_tasks.discard)
                
                # Pacing sleep
                elapsed = time.perf_counter() - start_time
                expected = len(running_tasks) * interval
                if elapsed < expected:
                    await asyncio.sleep(expected - elapsed)
            
            # Align to 1-second boundaries
            elapsed = time.perf_counter() - start_time
            if elapsed < 1.0:
                await asyncio.sleep(1.0 - elapsed)

if __name__ == "__main__":
    asyncio.run(main())