Arya Pub-Sub System

A high-performance, WebSocket-based pub-sub system with REST API built using FastAPI. Arya provides real-time message broadcasting with thread-safe in-memory storage, backpressure handling, and comprehensive error management.

🚀 Features

WebSocket-based pub-sub with real-time message delivery
REST API for topic management and system monitoring
Thread-safe in-memory storage with asyncio locks
Message history replay (configurable per subscription)
Backpressure handling with bounded message queues
Graceful error handling with proper error codes
Health monitoring and system statistics
Docker support for easy deployment
Comprehensive test suite with 90%+ coverage

🏗️ Architecture

┌─────────────────┐    ┌──────────────────┐    ┌─────────────────┐
│   WebSocket     │    │   PubSubManager  │    │   REST API      │
│   Handler       │◄──►│   (Thread-Safe)  │◄──►│   Endpoints     │
│                 │    │                  │    │                 │
│ • Connection    │    │ • Topic Storage  │    │ • Topic CRUD    │
│ • Message Route │    │ • Subscriptions  │    │ • Health Check  │
│ • Broadcasting  │    │ • Message Queue  │    │ • Statistics    │
└─────────────────┘    └──────────────────┘    └─────────────────┘

Core Components

PubSubManager: Thread-safe singleton managing topics, subscriptions, and message distribution
WebSocketManager: Handles individual WebSocket connections with heartbeat and error handling
Topic: Manages message history with configurable retention (default: 100 messages)
Message Queue: Bounded queues (100 messages) for backpressure handling

Project Structure

arya/
├── src/arya/                    # Main application code
│   ├── api/                     # REST API endpoints
│   ├── core/                    # Core business logic
│   │   ├── models.py           # Pydantic data models
│   │   └── pubsub.py           # PubSub manager
│   ├── websocket/              # WebSocket handlers
│   ├── config/                 # Configuration management
│   ├── exceptions/             # Custom exceptions
│   └── utils/                  # Utility functions
├── tests/                      # Test suite
├── simple_test.py              # Contract compliance test
├── test_client.py              # Interactive demo client
├── docker-compose.yml          # Docker setup
└── requirements.txt            # Dependencies

🚦 Quick Start

Using Docker (Recommended)

# Build and run with Docker
docker build -t arya-pubsub .
docker run -p 8000:8000 arya-pubsub

# Or use docker-compose
docker-compose up --build

Local Development

# Install dependencies
pip install -r requirements.txt

# Run the server
PYTHONPATH=src python src/main.py

# Or with uvicorn directly (development with auto-reload)
PYTHONPATH=src uvicorn arya.api.main:create_app --factory --host 0.0.0.0 --port 8000 --reload

The server will start on http://localhost:8000

📡 API Documentation

REST Endpoints

Create Topic

curl -X POST "http://localhost:8000/topics" \
     -H "Content-Type: application/json" \
     -d '{"name": "orders"}'

# Response: {"status": "created", "topic": "orders"}

Delete Topic

curl -X DELETE "http://localhost:8000/topics/orders"

# Response: {"status": "deleted", "topic": "orders"}

List Topics

curl "http://localhost:8000/topics"

# Response: {
#   "topics": [
#     {"name": "orders", "subscribers": 3, "messages": 42}
#   ]
# }

Health Check

curl "http://localhost:8000/health"

# Response: {
#   "uptime_sec": 3600,
#   "topics": 5,
#   "subscribers": 15
# }

System Statistics

curl "http://localhost:8000/stats"

# Response: {
#   "topics": {
#     "orders": {"messages": 42, "subscribers": 3},
#     "notifications": {"messages": 15, "subscribers": 7}
#   }
# }

WebSocket Protocol

Connect to WebSocket endpoint: ws://localhost:8000/ws

Message Format

All WebSocket messages use JSON format with the following structure:

{
  "type": "message_type",
  "request_id": "optional_correlation_id",
  // ... additional fields based on message type
}

Subscribe to Topic

{
  "type": "subscribe",
  "topic": "orders",
  "client_id": "unique-client-123",
  "last_n": 5,
  "request_id": "sub-1"
}

Response:

{
  "type": "ack",
  "message": "Subscribed to topic 'orders'",
  "request_id": "sub-1",
  "timestamp": "2024-01-15T10:30:00Z"
}

Unsubscribe from Topic

{
  "type": "unsubscribe",
  "topic": "orders",
  "client_id": "unique-client-123",
  "request_id": "unsub-1"
}

Publish Message

{
  "type": "publish",
  "topic": "orders",
  "message": {
    "order_id": 12345,
    "customer_id": "cust-789",
    "amount": 99.99,
    "items": ["item1", "item2"]
  },
  "request_id": "pub-1"
}

Receive Published Messages

{
  "type": "event",
  "topic": "orders",
  "message": {
    "id": "msg-uuid-here",
    "payload": {
      "order_id": 12345,
      "customer_id": "cust-789",
      "amount": 99.99
    },
    "timestamp": "2024-01-15T10:30:00Z"
  }
}

Heartbeat (Ping/Pong)

// Send ping
{"type": "ping", "request_id": "ping-1"}

// Receive pong
{"type": "pong", "request_id": "ping-1", "timestamp": "2024-01-15T10:30:00Z"}

Error Responses

{
  "type": "error",
  "error_code": "TOPIC_NOT_FOUND",
  "error_message": "Topic 'nonexistent' not found",
  "request_id": "sub-1",
  "timestamp": "2024-01-15T10:30:00Z"
}

Error Codes:

BAD_REQUEST: Invalid request format or parameters
TOPIC_NOT_FOUND: Requested topic doesn't exist
SLOW_CONSUMER: Client is too slow, connection terminated
INTERNAL_ERROR: Server-side error

🧪 Testing

Run Tests

# Create and activate virtual environment (if not already done)
python3 -m venv venv
source venv/bin/activate

# Install dependencies
pip install -r requirements.txt

# Run all tests
PYTHONPATH=src python -m pytest tests/ -v

# Run with coverage (if pytest-cov is installed)
PYTHONPATH=src python -m pytest tests/ --cov=src --cov-report=html

# Run specific test categories
PYTHONPATH=src python -m pytest tests/test_models.py tests/test_pubsub.py -v  # Unit tests
PYTHONPATH=src python -m pytest tests/test_api.py tests/test_websocket.py -v  # Integration tests

Test Categories

Unit Tests: test_models.py, test_pubsub.py
Integration Tests: test_api.py, test_websocket.py
Coverage Target: 90%+

Manual Testing

Basic Contract Compliance Test

# Start the service first (in one terminal):
source venv/bin/activate
PYTHONPATH=src python src/main.py

# Then run the test (in another terminal):
source venv/bin/activate
python simple_test.py

Interactive WebSocket Demo

# Start the service first (in one terminal):
source venv/bin/activate
PYTHONPATH=src python src/main.py

# Then run the test client (in another terminal):
source venv/bin/activate
python test_client.py

🔧 Configuration

Environment Variables

LOG_LEVEL: Logging level (DEBUG, INFO, WARNING, ERROR)
MAX_MESSAGE_HISTORY: Maximum messages per topic (default: 100)
MAX_QUEUE_SIZE: Maximum message queue size per connection (default: 100)
HEARTBEAT_INTERVAL: Heartbeat interval in seconds (default: 30)

Performance Tuning

Message History: Adjust per-topic message retention
Queue Size: Tune backpressure thresholds
Connection Limits: Configure via uvicorn workers
Memory Usage: Monitor with /stats endpoint

🚀 Production Deployment

Docker Production Setup

# Use multi-stage build for smaller image
FROM python:3.11-slim as builder
WORKDIR /app
COPY requirements.txt .
RUN pip install --user -r requirements.txt

FROM python:3.11-slim
WORKDIR /app
COPY --from=builder /root/.local /root/.local
COPY . .
RUN useradd --create-home --shell /bin/bash arya
USER arya
EXPOSE 8000
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "4"]

Kubernetes Deployment

apiVersion: apps/v1
kind: Deployment
metadata:
  name: arya-pubsub
spec:
  replicas: 3
  selector:
    matchLabels:
      app: arya-pubsub
  template:
    metadata:
      labels:
        app: arya-pubsub
    spec:
      containers:
      - name: arya-pubsub
        image: arya-pubsub:latest
        ports:
        - containerPort: 8000
        env:
        - name: LOG_LEVEL
          value: "INFO"
        resources:
          requests:
            memory: "256Mi"
            cpu: "250m"
          limits:
            memory: "512Mi"
            cpu: "500m"

Load Balancing

For horizontal scaling, use a load balancer with session affinity or implement client-side connection management for WebSocket persistence.

📊 Monitoring

Health Checks

# Basic health check
curl http://localhost:8000/health

# Detailed statistics
curl http://localhost:8000/stats

Metrics to Monitor

Connection Count: Active WebSocket connections
Topic Count: Number of active topics
Message Throughput: Messages per second
Memory Usage: Topic message history size
Error Rates: Failed connections and operations

Logging

The system provides structured logging with different levels:

# Critical operations
logger.info("Client subscribed to topic")

# Debug information  
logger.debug("Message added to topic history")

# Errors and warnings
logger.error("Failed to broadcast message")

🔒 Security Considerations

Current Implementation

Input Validation: Pydantic models validate all inputs
Error Handling: Sanitized error messages
Resource Limits: Bounded queues and message history
Non-root User: Docker runs as non-privileged user

Production Recommendations

Authentication: Add JWT or API key authentication
Rate Limiting: Implement per-client rate limits
TLS/SSL: Use HTTPS/WSS in production
Network Security: Deploy behind reverse proxy
Monitoring: Add security event logging

🚨 Limitations & Design Decisions

Current Limitations

In-Memory Only: No persistence across restarts
Single Instance: No clustering support
No Authentication: Open access (add auth layer)
Message Size: No explicit size limits (JSON parsing limits apply)

Design Decisions

Backpressure Policy: Disconnect slow consumers (prevents memory exhaustion)
Message History: 100 messages per topic (configurable)
Thread Safety: AsyncIO locks (not thread-safe across OS threads)
Error Handling: Fail fast with clear error messages

🛠️ Development

Project Structure

arya/
├── main.py              # FastAPI application
├── models.py            # Pydantic data models
├── pubsub.py           # Core pub-sub manager
├── websocket_handler.py # WebSocket connection management
├── requirements.txt     # Python dependencies
├── Dockerfile          # Container configuration
├── docker-compose.yml  # Local development setup
├── pytest.ini         # Test configuration
├── tests/              # Test suite
│   ├── conftest.py     # Test fixtures
│   ├── test_models.py  # Model validation tests
│   ├── test_pubsub.py  # Core logic tests
│   ├── test_api.py     # REST API tests
│   └── test_websocket.py # WebSocket tests
└── README.md           # This file

Adding Features

New Message Types: Add to models.py and handle in websocket_handler.py
New Endpoints: Add to main.py with proper error handling
Storage Backends: Extend PubSubManager with pluggable storage
Authentication: Add middleware to FastAPI application

Contributing

Fork the repository
Create feature branch
Add tests for new functionality
Ensure all tests pass
Submit pull request

📝 Example Usage

Python WebSocket Client

import asyncio
import websockets
import json

async def client_example():
    uri = "ws://localhost:8000/ws"
    
    async with websockets.connect(uri) as websocket:
        # Subscribe to topic
        subscribe_msg = {
            "type": "subscribe",
            "topic": "orders",
            "client_id": "python-client-1",
            "last_n": 5
        }
        await websocket.send(json.dumps(subscribe_msg))
        
        # Receive acknowledgment
        response = await websocket.recv()
        print(f"Subscribe response: {response}")
        
        # Publish a message
        publish_msg = {
            "type": "publish",
            "topic": "orders",
            "message": {"order_id": 123, "amount": 99.99}
        }
        await websocket.send(json.dumps(publish_msg))
        
        # Listen for events
        while True:
            message = await websocket.recv()
            data = json.loads(message)
            if data["type"] == "event":
                print(f"Received: {data['message']['payload']}")

# Run the client
asyncio.run(client_example())

JavaScript WebSocket Client

const ws = new WebSocket('ws://localhost:8000/ws');

ws.onopen = function() {
    // Subscribe to topic
    const subscribeMsg = {
        type: 'subscribe',
        topic: 'orders',
        client_id: 'js-client-1',
        last_n: 10
    };
    ws.send(JSON.stringify(subscribeMsg));
};

ws.onmessage = function(event) {
    const data = JSON.parse(event.data);
    
    if (data.type === 'event') {
        console.log('Received message:', data.message.payload);
    } else if (data.type === 'ack') {
        console.log('Operation acknowledged:', data.message);
    } else if (data.type === 'error') {
        console.error('Error:', data.error_message);
    }
};

// Publish message
function publishMessage(orderData) {
    const publishMsg = {
        type: 'publish',
        topic: 'orders',
        message: orderData
    };
    ws.send(JSON.stringify(publishMsg));
}

// Example usage
publishMessage({order_id: 456, customer: '[email protected]', amount: 149.99});

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

🤝 Support

For questions, issues, or contributions:

Create an issue on GitHub
Check existing documentation
Review test cases for usage examples

Built with ❤️ using FastAPI, WebSockets, and Python asyncio

Name		Name	Last commit message	Last commit date
Latest commit History 5 Commits
src		src
tests		tests
.gitignore		.gitignore
CHANGELOG.md		CHANGELOG.md
DEPLOYMENT.md		DEPLOYMENT.md
Dockerfile		Dockerfile
README.md		README.md
cursor-step-by-step-task-distribution.md		cursor-step-by-step-task-distribution.md
docker-compose.yml		docker-compose.yml
pytest.ini		pytest.ini
render.yaml		render.yaml
requirements-dev.txt		requirements-dev.txt
requirements.txt		requirements.txt
simple_test.py		simple_test.py
start.sh		start.sh
test_client.py		test_client.py

Uh oh!

Uh oh!

theiconik/arya

Folders and files

Latest commit

History

Repository files navigation

Arya Pub-Sub System

🚀 Features

🏗️ Architecture

Core Components

Project Structure

🚦 Quick Start

Using Docker (Recommended)

Local Development

📡 API Documentation

REST Endpoints

Create Topic

Delete Topic

List Topics

Health Check

System Statistics

WebSocket Protocol

Message Format

Subscribe to Topic

Unsubscribe from Topic

Publish Message

Receive Published Messages

Heartbeat (Ping/Pong)

Error Responses

🧪 Testing

Run Tests

Test Categories

Manual Testing

Basic Contract Compliance Test

Interactive WebSocket Demo

🔧 Configuration

Environment Variables

Performance Tuning

🚀 Production Deployment

Docker Production Setup

Kubernetes Deployment

Load Balancing

📊 Monitoring

Health Checks

Metrics to Monitor

Logging

🔒 Security Considerations

Current Implementation

Production Recommendations

🚨 Limitations & Design Decisions

Current Limitations

Design Decisions

🛠️ Development

Project Structure

Adding Features

Contributing

📝 Example Usage

Python WebSocket Client

JavaScript WebSocket Client

📄 License

🤝 Support

About

Topics

Resources

Uh oh!

Stars

Watchers

Forks

Packages 0

Languages

Packages