|
import asyncio
|
|
import time
|
|
import random
|
|
import statistics
|
|
from collections import deque
|
|
import websockets
|
|
|
|
# --- Test Configuration ---
|
|
HOST = "127.0.0.1"
|
|
PORT = 8080
|
|
URI = f"ws://{HOST}:{PORT}"
|
|
|
|
# Client setup
|
|
NUM_SUBSCRIBERS = 1000
|
|
NUM_PUBLISHERS = 10
|
|
CHANNELS = ["news", "sports", "tech", "finance", "weather"]
|
|
|
|
# Test execution
|
|
TEST_DURATION_S = 15
|
|
MESSAGES_PER_SECOND_PER_PUBLISHER = 100 # Increased message rate
|
|
|
|
# --- Global State & Metrics ---
|
|
latencies = deque()
|
|
messages_sent = 0
|
|
messages_received = 0
|
|
subscriber_setup_count = 0
|
|
all_subscribed_event = asyncio.Event()
|
|
|
|
async def subscriber_client(client_id: int):
|
|
global subscriber_setup_count, messages_received
|
|
channel = random.choice(CHANNELS)
|
|
|
|
try:
|
|
async with websockets.connect(URI) as websocket:
|
|
await websocket.send(f"sub {channel}")
|
|
subscriber_setup_count += 1
|
|
if subscriber_setup_count == NUM_SUBSCRIBERS:
|
|
print("✅ All subscribers are connected and subscribed. Starting publishers...")
|
|
all_subscribed_event.set()
|
|
|
|
while True:
|
|
message = await websocket.recv()
|
|
try:
|
|
sent_time_str = message.split(":", 1)[0]
|
|
sent_time = float(sent_time_str)
|
|
latency = time.time() - sent_time
|
|
latencies.append(latency)
|
|
messages_received += 1
|
|
except (ValueError, IndexError):
|
|
print(f"Warning: Received malformed message: {message}")
|
|
|
|
except (websockets.exceptions.ConnectionClosedError, ConnectionRefusedError) as e:
|
|
print(f"Subscriber {client_id} disconnected: {e}")
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as e:
|
|
print(f"An unexpected error occurred in subscriber {client_id}: {e}")
|
|
|
|
async def publisher_client(client_id: int):
|
|
global messages_sent
|
|
await all_subscribed_event.wait()
|
|
|
|
sleep_interval = 1.0 / MESSAGES_PER_SECOND_PER_PUBLISHER
|
|
|
|
try:
|
|
async with websockets.connect(URI) as websocket:
|
|
while True:
|
|
channel = random.choice(CHANNELS)
|
|
send_time = time.time()
|
|
message = f"{send_time:.6f}:Hello from publisher {client_id} on channel {channel}"
|
|
|
|
await websocket.send(f"pub {channel} {message}")
|
|
messages_sent += 1
|
|
|
|
await asyncio.sleep(sleep_interval)
|
|
|
|
except (websockets.exceptions.ConnectionClosedError, ConnectionRefusedError) as e:
|
|
print(f"Publisher {client_id} disconnected: {e}")
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as e:
|
|
print(f"An unexpected error occurred in publisher {client_id}: {e}")
|
|
|
|
def print_report():
|
|
print("\n" + "="*80)
|
|
print("PERFORMANCE REPORT".center(80))
|
|
print("="*80)
|
|
|
|
if not latencies:
|
|
print("No messages were received. Cannot generate a report. Is the server running?")
|
|
return
|
|
|
|
total_sent = messages_sent
|
|
total_received = messages_received
|
|
message_loss = max(0, total_sent - total_received)
|
|
loss_rate = (message_loss / total_sent * 100) if total_sent > 0 else 0
|
|
throughput = total_received / TEST_DURATION_S
|
|
|
|
print(f"Test Duration: {TEST_DURATION_S} seconds")
|
|
print(f"Total Messages Sent: {total_sent}")
|
|
print(f"Total Messages Rcvd: {total_received}")
|
|
print(f"Message Loss: {message_loss} ({loss_rate:.2f}%)")
|
|
print(f"Actual Throughput: {throughput:.2f} msg/sec")
|
|
print("-"*80)
|
|
|
|
sorted_latencies = sorted(latencies)
|
|
avg_latency_ms = statistics.mean(sorted_latencies) * 1000
|
|
min_latency_ms = sorted_latencies[0] * 1000
|
|
max_latency_ms = sorted_latencies[-1] * 1000
|
|
p50_latency_ms = statistics.median(sorted_latencies) * 1000
|
|
p95_latency_ms = sorted_latencies[int(len(sorted_latencies) * 0.95)] * 1000
|
|
p99_latency_ms = sorted_latencies[int(len(sorted_latencies) * 0.99)] * 1000
|
|
|
|
print("Latency Statistics (ms):")
|
|
print(f" Average: {avg_latency_ms:.4f} ms")
|
|
print(f" Min: {min_latency_ms:.4f} ms")
|
|
print(f" Max: {max_latency_ms:.4f} ms")
|
|
print(f" Median (p50): {p50_latency_ms:.4f} ms")
|
|
print(f" 95th Percentile: {p95_latency_ms:.4f} ms")
|
|
print(f" 99th Percentile: {p99_latency_ms:.4f} ms")
|
|
print("="*80)
|
|
|
|
async def main():
|
|
print("Starting WebSocket Pub/Sub Load Test...")
|
|
print(f"Simulating {NUM_SUBSCRIBERS} subscribers and {NUM_PUBLISHERS} publishers.")
|
|
print(f"Publishing at ~{NUM_PUBLISHERS * MESSAGES_PER_SECOND_PER_PUBLISHER} msg/sec for {TEST_DURATION_S} seconds.")
|
|
print("-"*80)
|
|
|
|
subscriber_tasks = [asyncio.create_task(subscriber_client(i)) for i in range(NUM_SUBSCRIBERS)]
|
|
publisher_tasks = [asyncio.create_task(publisher_client(i)) for i in range(NUM_PUBLISHERS)]
|
|
all_tasks = subscriber_tasks + publisher_tasks
|
|
|
|
try:
|
|
await asyncio.sleep(TEST_DURATION_S)
|
|
finally:
|
|
print("\nTest duration finished. Shutting down clients...")
|
|
for task in all_tasks:
|
|
task.cancel()
|
|
|
|
await asyncio.gather(*all_tasks, return_exceptions=True)
|
|
print_report()
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
print("\nTest interrupted by user.")
|
|
|