Websocket endpoint

.github/workflows/ci.yml

@@ -56,6 +56,10 @@ jobs:
       shell: bash -l {0}
       run: source continuous_integration/scripts/start_postgres.sh
 
+    - name: Start Redis service in container.
+      shell: bash -l {0}
+      run: source continuous_integration/scripts/start_redis.sh
+
 
     - name: Ensure example data is migrated to current catalog database schema.
       # The example data is expected to be kept up to date to the latest Tiled

.gitignore

@@ -82,3 +82,6 @@ pixi.lock
 # uv environments
 uv.lock
 .venv
+# pixi environments
+.pixi/*
+!.pixi/config.toml

continuous_integration/scripts/start_redis.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+set -e
+
+docker run -d --rm --name tiled-test-redis -p 6379:6379 docker.io/redis:7-alpine
+docker ps

locust/README.md

@@ -1,35 +1,107 @@
 # Tiled Load Testing with Locust
 
-Simple load testing for Tiled using the `reader.py` file.
+Load testing for Tiled using Locust. Two test files are available:
+- `reader.py` - Tests HTTP read operations and search endpoints
+- `streaming.py` - Tests streaming data writes and WebSocket delivery latency
 
 ## Quick Start
 
 ```bash
-# Install dependencies (dev environment includes locust)
-pixi install -e dev
+# Install dependencies (locust should already be available in the environment)
+# If not installed, add it to your requirements or install with:
+# uv add locust
 ```
 
+## Starting Test Server
+
+Before running locust tests, start a Tiled server:
+
+```bash
+# Basic server (works for most tests)
+uv run tiled serve catalog \
+  --host 0.0.0.0 \
+  --port 8000 \
+  --api-key secret \
+  --temp \
+  --init
+```
+
+For streaming tests with Redis cache (optional):
+```bash
+# Start Redis first
+redis-server
+
+# Start Tiled server with Redis cache
+uv run tiled serve catalog \
+  --host 0.0.0.0 \
+  --port 8000 \
+  --api-key secret \
+  --cache-uri "redis://localhost:6379" \
+  --cache-ttl 60 \
+  --temp \
+  --init
+```
+
+This creates a temporary catalog with:
+- API key authentication (key: "secret")
+- Temporary writable storage (automatically cleaned up)
+- Optional Redis cache for enhanced streaming performance
+- Server running on http://localhost:8000
+
+## Reading Performance Tests (`reader.py`)
+
+Tests various HTTP endpoints for reading data, metadata, and search operations.
+
 ### Examples
 Run with default localhost server (uses default API key 'secret'):
 ```bash
-pixi run -e dev locust -f reader.py --host http://localhost:8000
+uv run locust -f reader.py --headless -u 100 -r 10 -t 60s --host http://localhost:8000
 ```
 
 Run with custom API key:
 ```bash
-pixi run -e dev locust -f reader.py --host http://localhost:8000 --api-key your-api-key
+uv run locust -f reader.py --headless -u 100 -r 10 -t 60s --host http://localhost:8000 --api-key your-api-key
 ```
 
 Run with custom container name (defaults to locust_testing):
 ```bash
-pixi run -e dev locust -f reader.py --host http://localhost:8000 --container-name my_test_container
+uv run locust -f reader.py --headless -u 100 -r 10 -t 60s --host http://localhost:8000 --container-name my_test_container
+```
+
+## Streaming Performance Tests (`streaming.py`)
+
+Tests streaming data writes and WebSocket delivery with end-to-end latency measurement.
+
+**Note:** The `--node-name` parameter is required for streaming tests to avoid conflicts when multiple test runs create nodes with the same name.
+
+### Examples
+Run with required node name:
+```bash
+uv run locust -f streaming.py --headless -u 10 -r 2 -t 120s --host http://localhost:8000 --node-name my_test_stream
 ```
 
-## Headless Mode
-Run without the web interface:
+Run with custom API key:
 ```bash
-pixi run -e dev locust -f reader.py --headless -u 100 -r 10 -t 60s
+uv run locust -f streaming.py --headless -u 10 -r 2 -t 120s --host http://localhost:8000 --api-key your-api-key --node-name my_test_stream
 ```
-- `-u 100`: 100 concurrent users
-- `-r 10`: Spawn 10 users per second
-- `-t 60s`: Run for 60 seconds
+
+Control user types with environment variables:
+```bash
+# 2 writers for every 1 streaming reader
+WRITER_WEIGHT=2 STREAMING_WEIGHT=1 uv run locust -f streaming.py --headless -u 10 -r 2 -t 120s --host http://localhost:8000 --node-name my_test_stream
+```
+
+### Streaming Test Components
+- **WriterUser**: Writes timestamped array data to streaming nodes
+- **StreamingUser**: Connects via WebSocket to measure write-to-delivery latency
+
+## Parameters
+- `-u N`: N concurrent users
+- `-r N`: Spawn N users per second
+- `-t Ns`: Run for N seconds
+- `--headless`: Run without web interface (required for automation)
+
+## Notes
+- All examples use `--headless` mode for reliable automation
+- For streaming tests, `--node-name` is required to avoid conflicts
+- Use environment variables `WRITER_WEIGHT` and `STREAMING_WEIGHT` to control user distribution

locust/streaming.py

@@ -0,0 +1,238 @@
+import json
+import logging
+import os
+import threading
+import time
+from urllib.parse import urlparse
+
+import msgpack
+import numpy as np
+import websocket
+
+from locust import HttpUser, between, events, task
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@events.init_command_line_parser.add_listener
+def _(parser):
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default="secret",
+        help="API key for Tiled authentication (default: secret)",
+    )
+    parser.add_argument(
+        "--node-name",
+        type=str,
+        required=True,
+        help="Node name for streaming test (required)",
+    )
+
+
+@events.init.add_listener
+def on_locust_init(environment, **kwargs):
+    if environment.host is None:
+        raise ValueError(
+            "Host must be specified with --host argument, or through the web-ui."
+        )
+
+    # Create the streaming node once for all users
+    create_streaming_node(
+        environment.host,
+        environment.parsed_options.api_key,
+        environment.parsed_options.node_name,
+    )
+
+
+def create_streaming_node(host, api_key, node_name):
+    """Create a streaming array node using Tiled client"""
+    from tiled.client import from_uri
+
+    # Connect to Tiled server using client
+    client = from_uri(host, api_key=api_key)
+
+    # Create initial streaming array
+    arr = np.full(5, 0.0, dtype=np.float64)  # Initial array with zeros
+    client.write_array(arr, key=node_name)
+
+    logger.info(f"Created streaming node: {node_name}")
+    client.logout()
+
+
+class WriterUser(HttpUser):
+    """User that writes streaming data to a Tiled node"""
+
+    wait_time = between(0.1, 0.2)  # Wait 0.1-0.2 seconds between writes
+    weight = int(os.getenv("WRITER_WEIGHT", 1))
+
+    def on_start(self):
+        """Initialize user state"""
+        self.node_name = self.environment.parsed_options.node_name
+        self.message_count = 0
+        self.api_key = self.environment.parsed_options.api_key
+
+        # Set authentication header
+        self.client.headers.update({"Authorization": f"Apikey {self.api_key}"})
+
+    @task(10)  # Run 10x as often as cleanup
+    def write_data(self):
+        """Write streaming data to the node"""
+        # Create data with current timestamp as all values
+        current_time = time.time()
+        data = np.full(5, current_time, dtype=np.float64)
+        binary_data = data.tobytes()
+
+        # Post data to the streaming endpoint
+        response = self.client.put(
+            f"/api/v1/array/full/{self.node_name}",
+            data=binary_data,
+            headers={"Content-Type": "application/octet-stream"},
+        )
+
+        # Log status
+        if response.status_code == 200:
+            logger.debug(f"Wrote message {self.message_count} to node {self.node_name}")
+            self.message_count += 1
+        else:
+            logger.error(
+                f"Failed to write message {self.message_count}: {response.status_code} - {response.text}"
+            )
+
+    @task(1)
+    def cleanup(self):
+        """Periodically cleanup the stream"""
+        if self.message_count > 50:
+            # Close the stream
+            response = self.client.delete(f"/api/v1/stream/close/{self.node_name}")
+            if response.status_code == 200:
+                logger.info(f"Closed stream for node {self.node_name}")
+
+            # Reset message count (node persists for other users)
+            self.message_count = 0
+
+
+class StreamingUser(HttpUser):
+    """User that connects to websocket stream and measures latency"""
+
+    wait_time = between(1, 2)
+    weight = int(os.getenv("STREAMING_WEIGHT", 1))
+
+    def on_start(self):
+        """Connect to the streaming endpoint"""
+        self.node_name = self.environment.parsed_options.node_name
+        self.api_key = self.environment.parsed_options.api_key
+        self.envelope_format = "msgpack"  # Use msgpack for efficiency
+        self.ws = None
+        self.connected = False
+
+        # Set up authentication for HTTP requests
+        self.client.headers.update({"Authorization": f"Apikey {self.api_key}"})
+
+        self._connect_websocket()
+
+    def _connect_websocket(self):
+        """Connect to the websocket stream"""
+        try:
+            # Parse host to get websocket URL
+            parsed = urlparse(self.host)
+            ws_scheme = "wss" if parsed.scheme == "https" else "ws"
+            host = f"{ws_scheme}://{parsed.netloc}"
+
+            ws_url = f"{host}/api/v1/stream/single/{self.node_name}?envelope_format={self.envelope_format}&start=0"
+
+            # Create websocket connection
+            self.ws = websocket.WebSocketApp(
+                ws_url,
+                header=[
+                    f"Authorization: Apikey {self.api_key}"
+                ],  # Proper Apikey format for websockets
+                on_open=self._on_open,
+                on_message=self._on_message,
+                on_error=self._on_error,
+                on_close=self._on_close,
+            )
+
+            # Start websocket in background thread
+            self.ws_thread = threading.Thread(target=self.ws.run_forever)
+            self.ws_thread.daemon = True
+            self.ws_thread.start()
+
+            # Wait a bit for connection
+            time.sleep(0.5)
+
+        except Exception as e:
+            logger.error(f"Failed to connect websocket: {e}")
+
+    def _on_open(self, ws):
+        """Websocket connection opened"""
+        self.connected = True
+        logger.info(f"WebSocket connected to {self.node_name}")
+
+    def _on_message(self, ws, message):
+        """Process websocket messages and measure latency"""
+        try:
+            received_time = time.time()
+
+            if isinstance(message, bytes):
+                data = msgpack.unpackb(message)
+            else:
+                data = json.loads(message)
+
+            # Extract timestamp from the payload (first element of the array)
+            payload = data.get("payload")
+            if payload and len(payload) > 0:
+                # Convert bytes back to numpy array to get the timestamp
+                payload_array = np.frombuffer(payload, dtype=np.float64)
+                if len(payload_array) > 0:
+                    write_time = payload_array[0]
+                    latency_ms = (received_time - write_time) * 1000
+
+                    logger.debug(
+                        f"WS latency (sequence {data.get('sequence', 'N/A')}): {latency_ms:.1f}ms"
+                    )
+
+                    # Report to Locust
+                    events.request.fire(
+                        request_type="WS",
+                        name="write_to_websocket_delivery",
+                        response_time=latency_ms,
+                        response_length=len(message),
+                        exception=None,
+                    )
+
+        except Exception as e:
+            logger.error(f"Error processing message: {e}")
+            events.request.fire(
+                request_type="WS",
+                name="write_to_websocket_delivery",
+                response_time=0,
+                response_length=0,
+                exception=e,
+            )
+
+    def _on_error(self, ws, error):
+        """Websocket error occurred"""
+        logger.error(f"WebSocket error: {error}")
+        self.connected = False
+
+    def _on_close(self, ws, close_status_code, close_msg):
+        """Websocket connection closed"""
+        logger.info(f"WebSocket closed: {close_status_code} - {close_msg}")
+        self.connected = False
+
+    @task
+    def keep_alive(self):
+        """Dummy task to keep the user active while listening for messages"""
+        if not self.connected and self.ws:
+            # Try to reconnect if disconnected
+            logger.info("Attempting to reconnect WebSocket...")
+            self._connect_websocket()
+
+    def on_stop(self):
+        """Clean up websocket connection"""
+        if self.ws:
+            self.ws.close()
+            logger.info("WebSocket connection closed")

pixi.toml

@@ -35,6 +35,7 @@ entrypoints = "*"
 rich = "*"
 stamina = "*"
 watchfiles = "*"
+websockets = "*"
 
 [feature.compression.dependencies]
 # python-blosc2 = "*"  # not available on Python < 3.11
@@ -98,6 +99,7 @@ prometheus_client = "*"
 python-dateutil = "*"
 python-jose = "*"
 python-multipart = "*"
+redis-py = "*"
 sqlalchemy = ">=2"
 starlette = ">=0.38.0"
 uvicorn = "*"