diff --git a/.gitignore b/.gitignore index 17dd7eab..31d70e5b 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,4 @@ flamegraph.svg criterion/ bench/results/*.csv bench/results/*.json +.bench-venv/ diff --git a/Cargo.lock b/Cargo.lock index 885a472e..372c06be 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -727,6 +727,7 @@ version = "0.4.3" dependencies = [ "bytes", "ember-protocol", + "ember-server", "prost-reflect", "tempfile", "tokio", diff --git a/bench/README.md b/bench/README.md index bf2cdf85..548abec7 100644 --- a/bench/README.md +++ b/bench/README.md @@ -113,6 +113,30 @@ AES-256-GCM encryption at rest (AOF and snapshots). requires building with `--fe note: encryption only affects persistence writes. GET throughput should be unchanged since reads come from the in-memory keyspace. +### vector similarity + +ember vs chromadb vs pgvector. 100k random vectors, 128 dimensions, cosine metric, k=10 kNN search. +HNSW index: M=16, ef_construction=64 for all systems. tested on GCP c2-standard-8. + +| metric | ember | chromadb | pgvector | +|--------|-------|----------|----------| +| insert (vectors/sec) | 917 | **3,738** | 1,562 | +| query (queries/sec) | **1,214** | 376 | 882 | +| query p99 (ms) | **1.09ms** | 2.90ms | 1.52ms | +| memory (MB) | **30 MB** | 122 MB | 178 MB | + +ember's query throughput is 3.2x chromadb and 1.4x pgvector, with 4-6x lower memory usage. insert throughput is lower due to per-vector RESP protocol overhead — batched pipelining helps but each VADD is still a separate command. + +#### SIFT1M recall accuracy (128-dim, 1M vectors, 10k queries) + +| metric | ember | chromadb | pgvector | +|--------|-------|----------|----------| +| recall@10 | — | — | — | +| insert (vectors/sec) | — | — | — | +| query p99 (ms) | — | — | — | + +*results pending — requires a larger VM (c2-standard-16 or higher) since the 1M-vector HNSW index exceeds 16GB RAM during construction. run `bench/bench-vector.sh --sift` to populate.* + ### scaling efficiency | cores | ember sharded SET | scaling factor | @@ -166,6 +190,16 @@ cargo build --release -p ember-server --features jemalloc,encryption # comprehensive comparison using memtier_benchmark (redis + dragonfly) ./bench/bench-memtier.sh + +# vector similarity benchmark (requires --features vector, docker for full comparison) +cargo build --release -p ember-server --features jemalloc,vector +./bench/bench-vector.sh + +# vector benchmark (ember only, no docker required) +./bench/bench-vector.sh --ember-only + +# SIFT1M recall accuracy +./bench/bench-vector.sh --sift ``` ### cloud VM benchmarking @@ -204,7 +238,9 @@ gcloud compute instances delete ember-bench --zone=us-central1-a | `compare-redis.sh` | comprehensive comparison using redis-benchmark | | `bench-memtier.sh` | comprehensive comparison using memtier_benchmark | | `bench-encryption.sh` | encryption at rest overhead (plaintext vs AES-256-GCM) | +| `bench-vector.sh` | vector similarity: ember vs chromadb vs pgvector | | `setup-vm.sh` | bootstrap dependencies on fresh ubuntu VM | +| `setup-vm-vector.sh` | additional dependencies for vector benchmarks | ## configuration diff --git a/bench/bench-vector.py b/bench/bench-vector.py new file mode 100755 index 00000000..485d3f9e --- /dev/null +++ b/bench/bench-vector.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python3 +""" +vector similarity benchmark harness. + +benchmarks vector insert and kNN query performance across ember, chromadb, +and pgvector. called by bench-vector.sh with appropriate arguments. + +usage: + python3 bench/bench-vector.py --system ember --mode random --dim 128 --count 100000 + python3 bench/bench-vector.py --system chromadb --mode sift --sift-dir bench/vector_data +""" + +import argparse +import json +import time +import sys +import os +import numpy as np +from abc import ABC, abstractmethod + + +# --------------------------------------------------------------------------- +# vector generation +# --------------------------------------------------------------------------- + +def generate_vectors(count: int, dim: int, seed: int = 42) -> np.ndarray: + """generate random unit vectors for benchmarking.""" + rng = np.random.RandomState(seed) + vectors = rng.randn(count, dim).astype(np.float32) + # normalize to unit vectors for cosine similarity + norms = np.linalg.norm(vectors, axis=1, keepdims=True) + norms[norms == 0] = 1.0 + vectors /= norms + return vectors + + +# --------------------------------------------------------------------------- +# client abstractions +# --------------------------------------------------------------------------- + +class VectorClient(ABC): + """base class for vector database clients.""" + + @abstractmethod + def setup(self, dim: int, metric: str = "cosine"): + """create collection/index. called once before inserts.""" + + @abstractmethod + def insert_batch(self, ids: list, vectors: np.ndarray): + """insert a batch of vectors.""" + + @abstractmethod + def query(self, vector: np.ndarray, k: int) -> list: + """return top-k nearest neighbor ids.""" + + @abstractmethod + def teardown(self): + """clean up collection/index.""" + + @abstractmethod + def name(self) -> str: + """display name for this system.""" + + +class EmberClient(VectorClient): + """ember vector client using redis-py for RESP command execution.""" + + def __init__(self, host: str = "127.0.0.1", port: int = 6379): + import redis + self.conn = redis.Redis(host=host, port=port, decode_responses=True) + self.key = "bench_vectors" + + def setup(self, dim: int, metric: str = "cosine"): + # clear any previous data + self.conn.delete(self.key) + + def insert_batch(self, ids: list, vectors: np.ndarray): + pipe = self.conn.pipeline(transaction=False) + for i, vid in enumerate(ids): + vec = vectors[i] + # VADD key element v1 v2 ... METRIC COSINE M 16 EF 64 + args = [self.key, vid] + [str(float(v)) for v in vec] + args += ["METRIC", "COSINE", "M", "16", "EF", "64"] + pipe.execute_command("VADD", *args) + pipe.execute() + + def query(self, vector: np.ndarray, k: int) -> list: + args = [self.key] + [str(float(v)) for v in vector] + args += ["COUNT", str(k)] + result = self.conn.execute_command("VSIM", *args) + if result is None: + return [] + return [r.decode() if isinstance(r, bytes) else str(r) for r in result] + + def teardown(self): + self.conn.delete(self.key) + self.conn.close() + + def name(self) -> str: + return "ember" + + +class ChromaClient(VectorClient): + """chromadb client via HTTP API.""" + + def __init__(self, host: str = "127.0.0.1", port: int = 8000): + import chromadb + self.client = chromadb.HttpClient(host=host, port=port) + self.collection = None + self.collection_name = "bench_vectors" + + def setup(self, dim: int, metric: str = "cosine"): + # delete if exists + try: + self.client.delete_collection(self.collection_name) + except Exception: + pass + self.collection = self.client.create_collection( + name=self.collection_name, + metadata={ + "hnsw:space": "cosine", + "hnsw:M": 16, + "hnsw:construction_ef": 64, + }, + ) + + def insert_batch(self, ids: list, vectors: np.ndarray): + self.collection.add( + ids=ids, + embeddings=vectors.tolist(), + ) + + def query(self, vector: np.ndarray, k: int) -> list: + results = self.collection.query( + query_embeddings=[vector.tolist()], + n_results=k, + ) + return results["ids"][0] if results["ids"] else [] + + def teardown(self): + try: + self.client.delete_collection(self.collection_name) + except Exception: + pass + + def name(self) -> str: + return "chromadb" + + +class PgVectorClient(VectorClient): + """pgvector client via psycopg2.""" + + def __init__(self, host: str = "127.0.0.1", port: int = 5432, + user: str = "postgres", password: str = "postgres", + dbname: str = "vectordb"): + import psycopg2 + self.conn = psycopg2.connect( + host=host, port=port, user=user, password=password, dbname=dbname, + ) + self.conn.autocommit = True + self.dim = None + + def setup(self, dim: int, metric: str = "cosine"): + self.dim = dim + cur = self.conn.cursor() + cur.execute("CREATE EXTENSION IF NOT EXISTS vector") + cur.execute("DROP TABLE IF EXISTS bench_vectors") + cur.execute(f"CREATE TABLE bench_vectors (id text PRIMARY KEY, embedding vector({dim}))") + cur.close() + + def _create_index(self): + """create HNSW index after all inserts (faster than incremental).""" + cur = self.conn.cursor() + cur.execute( + "CREATE INDEX ON bench_vectors USING hnsw (embedding vector_cosine_ops) " + "WITH (m = 16, ef_construction = 64)" + ) + cur.close() + + def insert_batch(self, ids: list, vectors: np.ndarray): + cur = self.conn.cursor() + # format vectors as pgvector string literals: "[0.1,0.2,0.3]" + data = [ + (vid, "[" + ",".join(str(float(v)) for v in vec) + "]") + for vid, vec in zip(ids, vectors) + ] + cur.executemany( + "INSERT INTO bench_vectors (id, embedding) VALUES (%s, %s::vector) " + "ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding", + data, + ) + cur.close() + + def query(self, vector: np.ndarray, k: int) -> list: + cur = self.conn.cursor() + # explicit cast to vector type since psycopg2 sends lists as numeric[] + vec_str = "[" + ",".join(str(float(v)) for v in vector) + "]" + cur.execute( + "SELECT id FROM bench_vectors ORDER BY embedding <=> %s::vector LIMIT %s", + (vec_str, k), + ) + results = [row[0] for row in cur.fetchall()] + cur.close() + return results + + def teardown(self): + cur = self.conn.cursor() + cur.execute("DROP TABLE IF EXISTS bench_vectors") + cur.close() + self.conn.close() + + def name(self) -> str: + return "pgvector" + + +# --------------------------------------------------------------------------- +# benchmark functions +# --------------------------------------------------------------------------- + +def benchmark_insert(client: VectorClient, vectors: np.ndarray, + batch_size: int = 500) -> dict: + """measure insert throughput. returns vectors/sec.""" + n = len(vectors) + ids = [f"vec_{i}" for i in range(n)] + + start = time.perf_counter() + for i in range(0, n, batch_size): + end_idx = min(i + batch_size, n) + client.insert_batch(ids[i:end_idx], vectors[i:end_idx]) + elapsed = time.perf_counter() - start + + # for pgvector, create the HNSW index after all inserts + index_time = 0.0 + if isinstance(client, PgVectorClient): + idx_start = time.perf_counter() + client._create_index() + index_time = time.perf_counter() - idx_start + + throughput = n / elapsed if elapsed > 0 else 0 + + return { + "vectors": n, + "elapsed_sec": round(elapsed, 3), + "index_time_sec": round(index_time, 3), + "throughput": round(throughput, 1), + } + + +def benchmark_query(client: VectorClient, queries: np.ndarray, + k: int = 10, warmup: int = 50) -> dict: + """measure query latency and throughput.""" + n = len(queries) + + # warmup + for i in range(min(warmup, n)): + client.query(queries[i], k) + + # timed run — measure each query individually for latency percentiles + latencies = [] + for i in range(n): + start = time.perf_counter() + client.query(queries[i], k) + latencies.append(time.perf_counter() - start) + + latencies_ms = np.array(latencies) * 1000 + total_sec = sum(latencies) + throughput = n / total_sec if total_sec > 0 else 0 + + return { + "queries": n, + "elapsed_sec": round(total_sec, 3), + "throughput": round(throughput, 1), + "p50_ms": round(float(np.percentile(latencies_ms, 50)), 3), + "p95_ms": round(float(np.percentile(latencies_ms, 95)), 3), + "p99_ms": round(float(np.percentile(latencies_ms, 99)), 3), + } + + +def compute_recall(client: VectorClient, queries: np.ndarray, + ground_truth: np.ndarray, k: int = 10) -> dict: + """compute recall@k against ground truth nearest neighbors. + + ground_truth contains integer indices into the base vector array. + we use "vec_{index}" as the ID format matching our insert convention. + """ + n = len(queries) + hits = 0 + total = 0 + + for i in range(n): + predicted = set(client.query(queries[i], k)) + # ground truth gives integer indices; our IDs are "vec_{index}" + true_neighbors = set(f"vec_{gt}" for gt in ground_truth[i][:k]) + hits += len(predicted & true_neighbors) + total += k + + recall = hits / total if total > 0 else 0.0 + return { + "recall_at_k": round(recall, 4), + "k": k, + "queries": n, + } + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser(description="vector similarity benchmark") + parser.add_argument("--system", required=True, choices=["ember", "chromadb", "pgvector"]) + parser.add_argument("--mode", default="random", choices=["random", "sift"]) + parser.add_argument("--dim", type=int, default=128) + parser.add_argument("--count", type=int, default=100000) + parser.add_argument("--queries", type=int, default=1000) + parser.add_argument("--k", type=int, default=10) + parser.add_argument("--batch-size", type=int, default=500) + parser.add_argument("--ember-port", type=int, default=6379) + parser.add_argument("--chroma-port", type=int, default=8000) + parser.add_argument("--pgvector-port", type=int, default=5432) + parser.add_argument("--sift-dir", default="bench/vector_data") + parser.add_argument("--output", help="JSON output file") + args = parser.parse_args() + + # create client + if args.system == "ember": + client = EmberClient(port=args.ember_port) + elif args.system == "chromadb": + client = ChromaClient(port=args.chroma_port) + elif args.system == "pgvector": + client = PgVectorClient(port=args.pgvector_port) + + # load or generate vectors + if args.mode == "sift": + sys.path.insert(0, os.path.join(os.path.dirname(__file__), "vector_data")) + from sift_loader import load_sift1m + base_vectors, query_vectors, ground_truth = load_sift1m(args.sift_dir) + dim = base_vectors.shape[1] + print(f"loaded SIFT1M: {len(base_vectors)} base, {len(query_vectors)} queries, dim={dim}", file=sys.stderr) + else: + dim = args.dim + print(f"generating {args.count} random vectors (dim={dim})...", file=sys.stderr) + all_vectors = generate_vectors(args.count + args.queries, dim) + base_vectors = all_vectors[:args.count] + query_vectors = all_vectors[args.count:] + ground_truth = None + + # setup + client.setup(dim) + + # run benchmarks + print(f"benchmarking {client.name()} insert...", file=sys.stderr) + insert_result = benchmark_insert(client, base_vectors, batch_size=args.batch_size) + print(f" {insert_result['throughput']:.0f} vectors/sec", file=sys.stderr) + + print(f"benchmarking {client.name()} query (k={args.k})...", file=sys.stderr) + query_result = benchmark_query(client, query_vectors, k=args.k) + print(f" {query_result['throughput']:.0f} queries/sec, p99={query_result['p99_ms']:.2f}ms", file=sys.stderr) + + # recall (sift mode only) + recall_result = None + if args.mode == "sift" and ground_truth is not None: + print(f"computing recall@{args.k}...", file=sys.stderr) + recall_result = compute_recall(client, query_vectors, ground_truth, k=args.k) + print(f" recall@{args.k} = {recall_result['recall_at_k']:.4f}", file=sys.stderr) + + # cleanup + client.teardown() + + # output + result = { + "system": client.name(), + "mode": args.mode, + "dim": dim, + "insert": insert_result, + "query": query_result, + } + if recall_result: + result["recall"] = recall_result + + output_json = json.dumps(result, indent=2) + + if args.output: + with open(args.output, "w") as f: + f.write(output_json + "\n") + print(f"results written to {args.output}", file=sys.stderr) + else: + print(output_json) + + +if __name__ == "__main__": + main() diff --git a/bench/bench-vector.sh b/bench/bench-vector.sh new file mode 100755 index 00000000..75752ef3 --- /dev/null +++ b/bench/bench-vector.sh @@ -0,0 +1,482 @@ +#!/usr/bin/env bash +# +# vector similarity benchmark: ember vs chromadb vs pgvector. +# +# compares insert throughput, query throughput/latency, and memory usage +# across three vector databases using identical HNSW parameters (M=16, +# ef_construction=64) and cosine similarity. +# +# usage: +# bash bench/bench-vector.sh # full comparison (100k vectors) +# bash bench/bench-vector.sh --ember-only # ember only (no docker needed) +# bash bench/bench-vector.sh --quick # quick run (1k vectors) +# bash bench/bench-vector.sh --sift # SIFT1M recall accuracy +# +# environment variables: +# EMBER_PORT ember port (default: 6379) +# CHROMA_PORT chromadb port (default: 8000) +# PGVECTOR_PORT pgvector port (default: 5432) +# VECTOR_COUNT base vector count (default: 100000) +# VECTOR_DIM vector dimensions (default: 128) +# QUERY_COUNT query vector count (default: 1000) +# EMBER_BIN ember-server binary (default: ./target/release/ember-server) + +set -euo pipefail + +# --- configuration --- + +EMBER_PORT="${EMBER_PORT:-6379}" +CHROMA_PORT="${CHROMA_PORT:-8000}" +PGVECTOR_PORT="${PGVECTOR_PORT:-5432}" +VECTOR_COUNT="${VECTOR_COUNT:-100000}" +VECTOR_DIM="${VECTOR_DIM:-128}" +QUERY_COUNT="${QUERY_COUNT:-1000}" +K=10 +BATCH_SIZE=500 +EMBER_BIN="${EMBER_BIN:-./target/release/ember-server}" +RESULTS_DIR="bench/results" +TIMESTAMP=$(date +%Y%m%d-%H%M%S) +BENCH_SCRIPT="bench/bench-vector.py" +SIFT_DIR="bench/vector_data" + +EMBER_ONLY=false +QUICK_MODE=false +SIFT_MODE=false + +for arg in "$@"; do + case "$arg" in + --ember-only) EMBER_ONLY=true ;; + --quick) QUICK_MODE=true ;; + --sift) SIFT_MODE=true ;; + *) echo "unknown flag: $arg"; exit 1 ;; + esac +done + +if [[ "$QUICK_MODE" == "true" ]]; then + VECTOR_COUNT=1000 + QUERY_COUNT=100 +fi + +# --- cleanup --- + +EMBER_PID="" +CHROMA_CONTAINER="" +PGVECTOR_CONTAINER="" + +cleanup() { + echo "" + echo "cleaning up..." + [[ -n "$EMBER_PID" ]] && kill "$EMBER_PID" 2>/dev/null && wait "$EMBER_PID" 2>/dev/null || true + [[ -n "$CHROMA_CONTAINER" ]] && docker rm -f "$CHROMA_CONTAINER" > /dev/null 2>&1 || true + [[ -n "$PGVECTOR_CONTAINER" ]] && docker rm -f "$PGVECTOR_CONTAINER" > /dev/null 2>&1 || true + rm -f /tmp/bench_vector_*.json +} +trap cleanup EXIT + +# --- helpers --- + +format_number() { + printf "%'d" "$1" +} + +format_float() { + printf "%.1f" "$1" +} + +get_rss_mb() { + local pid=$1 + if [[ "$(uname)" == "Darwin" ]]; then + # macOS: ps reports RSS in bytes + local rss_bytes + rss_bytes=$(ps -o rss= -p "$pid" 2>/dev/null | tr -d ' ') + if [[ -n "$rss_bytes" ]]; then + echo $(( rss_bytes / 1024 )) + else + echo 0 + fi + else + # linux: /proc gives RSS in kB + local rss_kb + rss_kb=$(awk '/VmRSS/{print $2}' "/proc/$pid/status" 2>/dev/null || echo 0) + echo $(( rss_kb / 1024 )) + fi +} + +get_container_rss_mb() { + local container=$1 + # docker stats gives memory usage directly + local mem + mem=$(docker stats --no-stream --format "{{.MemUsage}}" "$container" 2>/dev/null | awk '{print $1}') + if [[ "$mem" == *GiB* ]]; then + local val=${mem%GiB*} + echo "$val" | awk '{printf "%d", $1 * 1024}' + elif [[ "$mem" == *MiB* ]]; then + local val=${mem%MiB*} + echo "$val" | awk '{printf "%d", $1}' + else + echo 0 + fi +} + +wait_for_ember() { + local port=$1 + local retries=50 + while ! redis-cli -p "$port" ping > /dev/null 2>&1; do + retries=$((retries - 1)) + if [[ $retries -le 0 ]]; then + echo "error: ember did not start on port $port" >&2 + exit 1 + fi + sleep 0.1 + done +} + +wait_for_chroma() { + local port=$1 + local retries=60 + while ! curl -sf "http://127.0.0.1:$port/api/v2/heartbeat" > /dev/null 2>&1; do + retries=$((retries - 1)) + if [[ $retries -le 0 ]]; then + echo "error: chromadb did not start on port $port" >&2 + exit 1 + fi + sleep 0.5 + done +} + +wait_for_pgvector() { + local container=$1 + local retries=60 + while ! docker exec "$container" pg_isready -U postgres > /dev/null 2>&1; do + retries=$((retries - 1)) + if [[ $retries -le 0 ]]; then + echo "error: pgvector did not start" >&2 + exit 1 + fi + sleep 0.5 + done + # extra pause for pgvector to fully initialize + sleep 1 +} + +# --- dependency checks --- + +if ! command -v python3 &> /dev/null; then + echo "error: python3 required" >&2 + exit 1 +fi + +if ! command -v redis-cli &> /dev/null; then + echo "error: redis-cli required (for ember health checks)" >&2 + exit 1 +fi + +# set up python venv for benchmark dependencies. +# this avoids issues with externally-managed environments (PEP 668) +# on macOS homebrew and modern linux distros. +VENV_DIR=".bench-venv" + +ensure_venv() { + if [[ ! -d "$VENV_DIR" ]]; then + echo "creating python venv for benchmark dependencies..." + python3 -m venv "$VENV_DIR" + fi + # shellcheck disable=SC1091 + source "$VENV_DIR/bin/activate" +} + +REQUIRED_DEPS="redis numpy" +if [[ "$EMBER_ONLY" == "false" ]]; then + REQUIRED_DEPS="redis numpy chromadb psycopg2-binary" + + if ! command -v docker &> /dev/null; then + echo "error: docker required for chromadb/pgvector (or use --ember-only)" >&2 + exit 1 + fi +fi + +# check if all deps are available in current python; if not, use venv +if ! python3 -c "import redis, numpy" 2>/dev/null; then + ensure_venv + pip install --quiet $REQUIRED_DEPS +elif [[ "$EMBER_ONLY" == "false" ]] && ! python3 -c "import chromadb, psycopg2" 2>/dev/null; then + ensure_venv + pip install --quiet $REQUIRED_DEPS +fi + +# build ember if needed +if [[ ! -x "$EMBER_BIN" ]]; then + echo "building ember-server with vector support..." + cargo build --release -p ember-server --features jemalloc,vector +fi + +# download SIFT dataset if needed +if [[ "$SIFT_MODE" == "true" ]]; then + if [[ ! -f "$SIFT_DIR/sift/sift_base.fvecs" ]]; then + echo "downloading SIFT1M dataset..." + bash bench/vector_data/download_sift.sh "$SIFT_DIR" + fi +fi + +mkdir -p "$RESULTS_DIR" + +# --- print configuration --- + +MODE="random" +if [[ "$SIFT_MODE" == "true" ]]; then + MODE="sift" +fi + +echo "" +echo "=== vector benchmark configuration ===" +echo "mode: $MODE" +if [[ "$SIFT_MODE" == "true" ]]; then + echo "dataset: SIFT1M (1M base, 10k queries, 128-dim)" +else + echo "vectors: $VECTOR_COUNT base, $QUERY_COUNT queries" + echo "dimensions: $VECTOR_DIM" +fi +echo "k: $K" +echo "hnsw: M=16, ef_construction=64" +echo "metric: cosine" +echo "" + +# --- common benchmark args --- + +BENCH_ARGS=( + --dim "$VECTOR_DIM" + --count "$VECTOR_COUNT" + --queries "$QUERY_COUNT" + --k "$K" + --batch-size "$BATCH_SIZE" + --ember-port "$EMBER_PORT" + --chroma-port "$CHROMA_PORT" + --pgvector-port "$PGVECTOR_PORT" + --sift-dir "$SIFT_DIR" +) + +if [[ "$SIFT_MODE" == "true" ]]; then + BENCH_ARGS+=(--mode sift) +else + BENCH_ARGS+=(--mode random) +fi + +# --- benchmark ember --- + +echo "--- ember ---" +echo "" +echo " starting ember on port $EMBER_PORT..." +"$EMBER_BIN" --port "$EMBER_PORT" > /dev/null 2>&1 & +EMBER_PID=$! +wait_for_ember "$EMBER_PORT" + +EMBER_JSON="/tmp/bench_vector_ember.json" +python3 "$BENCH_SCRIPT" --system ember "${BENCH_ARGS[@]}" --output "$EMBER_JSON" + +EMBER_RSS=$(get_rss_mb "$EMBER_PID") +echo " memory (RSS): ${EMBER_RSS} MB" + +kill "$EMBER_PID" 2>/dev/null && wait "$EMBER_PID" 2>/dev/null || true +EMBER_PID="" +sleep 0.5 +echo "" + +# --- benchmark chromadb --- + +CHROMA_JSON="/tmp/bench_vector_chromadb.json" +CHROMA_RSS=0 + +if [[ "$EMBER_ONLY" == "false" ]]; then + echo "--- chromadb ---" + echo "" + echo " starting chromadb container on port $CHROMA_PORT..." + CHROMA_CONTAINER=$(docker run -d --rm -p "$CHROMA_PORT:8000" chromadb/chroma 2>/dev/null) + wait_for_chroma "$CHROMA_PORT" + + python3 "$BENCH_SCRIPT" --system chromadb "${BENCH_ARGS[@]}" --output "$CHROMA_JSON" + + CHROMA_RSS=$(get_container_rss_mb "$CHROMA_CONTAINER") + echo " memory (RSS): ${CHROMA_RSS} MB" + + docker rm -f "$CHROMA_CONTAINER" > /dev/null 2>&1 || true + CHROMA_CONTAINER="" + sleep 0.5 + echo "" +fi + +# --- benchmark pgvector --- + +PGVECTOR_JSON="/tmp/bench_vector_pgvector.json" +PGVECTOR_RSS=0 + +if [[ "$EMBER_ONLY" == "false" ]]; then + echo "--- pgvector ---" + echo "" + echo " starting pgvector container on port $PGVECTOR_PORT..." + PGVECTOR_CONTAINER=$(docker run -d --rm \ + --shm-size=512m \ + -p "$PGVECTOR_PORT:5432" \ + -e POSTGRES_PASSWORD=postgres \ + -e POSTGRES_DB=vectordb \ + -e POSTGRES_HOST_AUTH_METHOD=trust \ + pgvector/pgvector:pg17 \ + postgres -c shared_buffers=256MB -c maintenance_work_mem=256MB \ + 2>/dev/null) + wait_for_pgvector "$PGVECTOR_CONTAINER" + + python3 "$BENCH_SCRIPT" --system pgvector "${BENCH_ARGS[@]}" --output "$PGVECTOR_JSON" + + PGVECTOR_RSS=$(get_container_rss_mb "$PGVECTOR_CONTAINER") + echo " memory (RSS): ${PGVECTOR_RSS} MB" + + docker rm -f "$PGVECTOR_CONTAINER" > /dev/null 2>&1 || true + PGVECTOR_CONTAINER="" + sleep 0.5 + echo "" +fi + +# --- parse results --- + +# extract values from JSON files using python (avoids jq dependency) +extract() { + local file=$1 + local key=$2 + python3 -c " +import json, sys +try: + d = json.load(open('$file')) + keys = '$key'.split('.') + v = d + for k in keys: + v = v[k] + print(v) +except Exception: + print('—') +" 2>/dev/null +} + +# ember results (always present) +E_INSERT=$(extract "$EMBER_JSON" "insert.throughput") +E_QUERY=$(extract "$EMBER_JSON" "query.throughput") +E_P50=$(extract "$EMBER_JSON" "query.p50_ms") +E_P95=$(extract "$EMBER_JSON" "query.p95_ms") +E_P99=$(extract "$EMBER_JSON" "query.p99_ms") + +# chromadb results +C_INSERT="—" +C_QUERY="—" +C_P50="—" +C_P95="—" +C_P99="—" + +if [[ "$EMBER_ONLY" == "false" ]] && [[ -f "$CHROMA_JSON" ]]; then + C_INSERT=$(extract "$CHROMA_JSON" "insert.throughput") + C_QUERY=$(extract "$CHROMA_JSON" "query.throughput") + C_P50=$(extract "$CHROMA_JSON" "query.p50_ms") + C_P95=$(extract "$CHROMA_JSON" "query.p95_ms") + C_P99=$(extract "$CHROMA_JSON" "query.p99_ms") +fi + +# pgvector results +P_INSERT="—" +P_QUERY="—" +P_P50="—" +P_P95="—" +P_P99="—" + +if [[ "$EMBER_ONLY" == "false" ]] && [[ -f "$PGVECTOR_JSON" ]]; then + P_INSERT=$(extract "$PGVECTOR_JSON" "insert.throughput") + P_QUERY=$(extract "$PGVECTOR_JSON" "query.throughput") + P_P50=$(extract "$PGVECTOR_JSON" "query.p50_ms") + P_P95=$(extract "$PGVECTOR_JSON" "query.p95_ms") + P_P99=$(extract "$PGVECTOR_JSON" "query.p99_ms") +fi + +# --- display results --- + +DATE=$(date +%Y-%m-%d) +CPU_CORES=$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 1) + +echo "" +echo "========================================================================" +echo " vector benchmark results — $DATE" +echo "========================================================================" +echo "" +echo "system: $CPU_CORES cores | $MODE mode" +if [[ "$SIFT_MODE" == "true" ]]; then + echo "dataset: SIFT1M (1M vectors, 128-dim, 10k queries)" +else + echo "vectors: $VECTOR_COUNT base, $QUERY_COUNT queries, ${VECTOR_DIM}-dim" +fi +echo "hnsw: M=16, ef_construction=64, cosine similarity, k=$K" +echo "" + +if [[ "$EMBER_ONLY" == "true" ]]; then + fmt="%-24s %14s" + printf "$fmt\n" "metric" "ember" + printf "$fmt\n" "------" "-----" + printf "$fmt\n" "insert (vectors/sec)" "$E_INSERT" + printf "$fmt\n" "query (queries/sec)" "$E_QUERY" + printf "$fmt\n" "query p50 (ms)" "$E_P50" + printf "$fmt\n" "query p95 (ms)" "$E_P95" + printf "$fmt\n" "query p99 (ms)" "$E_P99" + printf "$fmt\n" "memory (MB)" "$EMBER_RSS" +else + fmt="%-24s %14s %14s %14s" + printf "$fmt\n" "metric" "ember" "chromadb" "pgvector" + printf "$fmt\n" "------" "-----" "--------" "--------" + printf "$fmt\n" "insert (vectors/sec)" "$E_INSERT" "$C_INSERT" "$P_INSERT" + printf "$fmt\n" "query (queries/sec)" "$E_QUERY" "$C_QUERY" "$P_QUERY" + printf "$fmt\n" "query p50 (ms)" "$E_P50" "$C_P50" "$P_P50" + printf "$fmt\n" "query p95 (ms)" "$E_P95" "$C_P95" "$P_P95" + printf "$fmt\n" "query p99 (ms)" "$E_P99" "$C_P99" "$P_P99" + printf "$fmt\n" "memory (MB)" "$EMBER_RSS" "$CHROMA_RSS" "$PGVECTOR_RSS" +fi + +# recall (sift mode only) +if [[ "$SIFT_MODE" == "true" ]]; then + echo "" + echo "=== recall@$K ===" + echo "" + + E_RECALL=$(extract "$EMBER_JSON" "recall.recall_at_k") + + if [[ "$EMBER_ONLY" == "true" ]]; then + printf "%-24s %14s\n" "metric" "ember" + printf "%-24s %14s\n" "------" "-----" + printf "%-24s %14s\n" "recall@$K" "$E_RECALL" + else + C_RECALL=$(extract "$CHROMA_JSON" "recall.recall_at_k") + P_RECALL=$(extract "$PGVECTOR_JSON" "recall.recall_at_k") + printf "%-24s %14s %14s %14s\n" "metric" "ember" "chromadb" "pgvector" + printf "%-24s %14s %14s %14s\n" "------" "-----" "--------" "--------" + printf "%-24s %14s %14s %14s\n" "recall@$K" "$E_RECALL" "$C_RECALL" "$P_RECALL" + fi +fi + +echo "" + +# --- save CSV --- + +RESULT_FILE="$RESULTS_DIR/${TIMESTAMP}-vector.csv" +{ + echo "metric,ember,chromadb,pgvector" + echo "insert_throughput,$E_INSERT,$C_INSERT,$P_INSERT" + echo "query_throughput,$E_QUERY,$C_QUERY,$P_QUERY" + echo "query_p50_ms,$E_P50,$C_P50,$P_P50" + echo "query_p95_ms,$E_P95,$C_P95,$P_P95" + echo "query_p99_ms,$E_P99,$C_P99,$P_P99" + echo "memory_mb,$EMBER_RSS,$CHROMA_RSS,$PGVECTOR_RSS" + if [[ "$SIFT_MODE" == "true" ]]; then + E_RECALL=$(extract "$EMBER_JSON" "recall.recall_at_k") + C_RECALL="—" + P_RECALL="—" + if [[ "$EMBER_ONLY" == "false" ]]; then + C_RECALL=$(extract "$CHROMA_JSON" "recall.recall_at_k") + P_RECALL=$(extract "$PGVECTOR_JSON" "recall.recall_at_k") + fi + echo "recall_at_$K,$E_RECALL,$C_RECALL,$P_RECALL" + fi +} > "$RESULT_FILE" + +echo "raw results saved to $RESULT_FILE" diff --git a/bench/setup-vm-vector.sh b/bench/setup-vm-vector.sh new file mode 100755 index 00000000..04c187af --- /dev/null +++ b/bench/setup-vm-vector.sh @@ -0,0 +1,58 @@ +#!/bin/bash +# +# additional VM setup for vector similarity benchmarks. +# run after setup-vm.sh on a fresh ubuntu VM. +# +# installs: python3 + deps, docker, chromadb/pgvector images +# +# usage: ssh user@vm 'bash -s' < ./bench/setup-vm-vector.sh + +set -e + +echo "=== installing python dependencies ===" + +sudo apt-get update +sudo apt-get install -y python3-pip python3-dev + +pip3 install --quiet redis chromadb psycopg2-binary numpy + +echo "" +echo "=== installing docker ===" + +if command -v docker &> /dev/null; then + echo "docker already installed" + docker --version +else + curl -fsSL https://get.docker.com | sh + sudo usermod -aG docker "$USER" + echo "docker installed (you may need to log out and back in for group permissions)" +fi + +echo "" +echo "=== pulling docker images ===" + +sudo docker pull chromadb/chroma +sudo docker pull pgvector/pgvector:pg17 + +echo "" +echo "=== rebuilding ember with vector support ===" + +cd ~/ember +source ~/.cargo/env +cargo build --release -p ember-server --features jemalloc,vector + +echo "" +echo "=== verifying ===" + +python3 -c "import redis, chromadb, psycopg2, numpy; print('python deps: ok')" +sudo docker images | grep -E "chroma|pgvector" || true +./target/release/ember-server --help | head -3 + +echo "" +echo "=== vector benchmark setup complete ===" +echo "" +echo "run benchmarks with:" +echo " ./bench/bench-vector.sh # 100k random vectors" +echo " ./bench/bench-vector.sh --ember-only # ember only (no docker)" +echo " ./bench/bench-vector.sh --quick # quick sanity check" +echo " ./bench/bench-vector.sh --sift # SIFT1M recall accuracy" diff --git a/bench/vector_data/.gitignore b/bench/vector_data/.gitignore new file mode 100644 index 00000000..195c3d09 --- /dev/null +++ b/bench/vector_data/.gitignore @@ -0,0 +1,4 @@ +*.fvecs +*.ivecs +*.tar.gz +sift/ diff --git a/bench/vector_data/download_sift.sh b/bench/vector_data/download_sift.sh new file mode 100755 index 00000000..71c7a55e --- /dev/null +++ b/bench/vector_data/download_sift.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +# +# download SIFT1M dataset for recall benchmarks. +# +# source: ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz +# 1M base vectors (128-dim), 10k queries with ground truth. +# +# usage: bash bench/vector_data/download_sift.sh [output_dir] + +set -euo pipefail + +DATA_DIR="${1:-$(dirname "$0")}" + +if [[ -f "$DATA_DIR/sift/sift_base.fvecs" ]]; then + echo "SIFT1M already downloaded in $DATA_DIR/sift/" + exit 0 +fi + +echo "downloading SIFT1M dataset (~160MB)..." +curl -fSL "ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz" -o "$DATA_DIR/sift.tar.gz" + +echo "extracting..." +tar -xzf "$DATA_DIR/sift.tar.gz" -C "$DATA_DIR" +rm -f "$DATA_DIR/sift.tar.gz" + +# verify key files exist and have expected sizes +for f in sift_base.fvecs sift_query.fvecs sift_groundtruth.ivecs; do + if [[ ! -f "$DATA_DIR/sift/$f" ]]; then + echo "error: expected file $f not found after extraction" >&2 + exit 1 + fi +done + +echo "SIFT1M dataset ready in $DATA_DIR/sift/" +echo " base vectors: $(wc -c < "$DATA_DIR/sift/sift_base.fvecs") bytes" +echo " query vectors: $(wc -c < "$DATA_DIR/sift/sift_query.fvecs") bytes" +echo " ground truth: $(wc -c < "$DATA_DIR/sift/sift_groundtruth.ivecs") bytes" diff --git a/bench/vector_data/sift_loader.py b/bench/vector_data/sift_loader.py new file mode 100644 index 00000000..09a41130 --- /dev/null +++ b/bench/vector_data/sift_loader.py @@ -0,0 +1,43 @@ +""" +SIFT1M dataset loader. + +reads .fvecs and .ivecs binary formats used by the texmex corpus. +format: each vector is prefixed by a 4-byte int32 dimension, followed by +dim float32 (fvecs) or int32 (ivecs) values. +""" + +import numpy as np +import os + + +def read_fvecs(filename: str) -> np.ndarray: + """read vectors from .fvecs file format.""" + with open(filename, "rb") as f: + data = np.fromfile(f, dtype=np.float32) + dim = int(data[0].view(np.int32)) + # each vector is (1 + dim) float32s: dimension prefix + values + return data.reshape(-1, dim + 1)[:, 1:].copy() + + +def read_ivecs(filename: str) -> np.ndarray: + """read vectors from .ivecs file format.""" + with open(filename, "rb") as f: + data = np.fromfile(f, dtype=np.int32) + dim = int(data[0]) + return data.reshape(-1, dim + 1)[:, 1:].copy() + + +def load_sift1m(data_dir: str): + """ + load the SIFT1M dataset. + + returns: + base_vectors: (1000000, 128) float32 array + query_vectors: (10000, 128) float32 array + ground_truth: (10000, 100) int32 array of nearest neighbor indices + """ + sift_dir = os.path.join(data_dir, "sift") + base = read_fvecs(os.path.join(sift_dir, "sift_base.fvecs")) + queries = read_fvecs(os.path.join(sift_dir, "sift_query.fvecs")) + ground_truth = read_ivecs(os.path.join(sift_dir, "sift_groundtruth.ivecs")) + return base, queries, ground_truth diff --git a/crates/ember-core/src/concurrent.rs b/crates/ember-core/src/concurrent.rs index d0e3a6d4..cec23a87 100644 --- a/crates/ember-core/src/concurrent.rs +++ b/crates/ember-core/src/concurrent.rs @@ -726,7 +726,11 @@ mod tests { #[test] fn persist_removes_ttl() { let ks = ConcurrentKeyspace::default(); - ks.set("key".into(), Bytes::from("val"), Some(Duration::from_secs(60))); + ks.set( + "key".into(), + Bytes::from("val"), + Some(Duration::from_secs(60)), + ); assert!(ks.persist("key")); assert!(matches!(ks.ttl("key"), TtlResult::NoExpiry)); } diff --git a/crates/ember-server/src/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs index 81122eb9..ffdfe3de 100644 --- a/crates/ember-server/src/concurrent_handler.rs +++ b/crates/ember-server/src/concurrent_handler.rs @@ -266,14 +266,13 @@ async fn execute_concurrent( Command::Persist { key } => Frame::Integer(if keyspace.persist(&key) { 1 } else { 0 }), - Command::Pexpire { - key, - milliseconds, - } => Frame::Integer(if keyspace.pexpire(&key, milliseconds) { - 1 - } else { - 0 - }), + Command::Pexpire { key, milliseconds } => { + Frame::Integer(if keyspace.pexpire(&key, milliseconds) { + 1 + } else { + 0 + }) + } Command::Pttl { key } => match keyspace.pttl(&key) { TtlResult::Milliseconds(ms) => Frame::Integer(ms as i64), @@ -351,9 +350,7 @@ async fn execute_concurrent( Err(msg) => Frame::Error(msg.into()), }, - Command::Info { section } => { - render_concurrent_info(keyspace, ctx, section.as_deref()) - } + Command::Info { section } => render_concurrent_info(keyspace, ctx, section.as_deref()), // -- pub/sub -- Command::Publish { channel, message } => { diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index 41529f9f..7a8a48a1 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -713,8 +713,13 @@ async fn dispatch_command( ($key:expr, $req:expr, $tag:expr) => {{ let idx = engine.shard_for_key(&$key); match engine.dispatch_to_shard(idx, $req).await { - Ok(rx) => PendingResponse::Pending { rx, tag: $tag, start, cmd_name }, - Err(e) => PendingResponse::Immediate(Frame::Error(format!("ERR {e}"))) + Ok(rx) => PendingResponse::Pending { + rx, + tag: $tag, + start, + cmd_name, + }, + Err(e) => PendingResponse::Immediate(Frame::Error(format!("ERR {e}"))), } }}; } @@ -729,12 +734,28 @@ async fn dispatch_command( Command::Get { key } => { dispatch!(key, ShardRequest::Get { key }, ResponseTag::Get) } - Command::Set { key, value, expire, nx, xx } => { + Command::Set { + key, + value, + expire, + nx, + xx, + } => { let duration = expire.map(|e| match e { SetExpire::Ex(secs) => Duration::from_secs(secs), SetExpire::Px(millis) => Duration::from_millis(millis), }); - dispatch!(key, ShardRequest::Set { key, value, expire: duration, nx, xx }, ResponseTag::Set) + dispatch!( + key, + ShardRequest::Set { + key, + value, + expire: duration, + nx, + xx + }, + ResponseTag::Set + ) } Command::Incr { key } => { dispatch!(key, ShardRequest::Incr { key }, ResponseTag::IntResult) @@ -743,22 +764,42 @@ async fn dispatch_command( dispatch!(key, ShardRequest::Decr { key }, ResponseTag::IntResult) } Command::IncrBy { key, delta } => { - dispatch!(key, ShardRequest::IncrBy { key, delta }, ResponseTag::IntResult) + dispatch!( + key, + ShardRequest::IncrBy { key, delta }, + ResponseTag::IntResult + ) } Command::DecrBy { key, delta } => { - dispatch!(key, ShardRequest::DecrBy { key, delta }, ResponseTag::IntResult) + dispatch!( + key, + ShardRequest::DecrBy { key, delta }, + ResponseTag::IntResult + ) } Command::IncrByFloat { key, delta } => { - dispatch!(key, ShardRequest::IncrByFloat { key, delta }, ResponseTag::FloatResult) + dispatch!( + key, + ShardRequest::IncrByFloat { key, delta }, + ResponseTag::FloatResult + ) } Command::Append { key, value } => { - dispatch!(key, ShardRequest::Append { key, value }, ResponseTag::LenResultOom) + dispatch!( + key, + ShardRequest::Append { key, value }, + ResponseTag::LenResultOom + ) } Command::Strlen { key } => { dispatch!(key, ShardRequest::Strlen { key }, ResponseTag::LenResult) } Command::Expire { key, seconds } => { - dispatch!(key, ShardRequest::Expire { key, seconds }, ResponseTag::BoolToInt) + dispatch!( + key, + ShardRequest::Expire { key, seconds }, + ResponseTag::BoolToInt + ) } Command::Ttl { key } => { dispatch!(key, ShardRequest::Ttl { key }, ResponseTag::Ttl) @@ -770,7 +811,11 @@ async fn dispatch_command( dispatch!(key, ShardRequest::Pttl { key }, ResponseTag::Pttl) } Command::Pexpire { key, milliseconds } => { - dispatch!(key, ShardRequest::Pexpire { key, milliseconds }, ResponseTag::BoolToInt) + dispatch!( + key, + ShardRequest::Pexpire { key, milliseconds }, + ResponseTag::BoolToInt + ) } Command::Type { key } => { dispatch!(key, ShardRequest::Type { key }, ResponseTag::TypeResult) @@ -778,10 +823,18 @@ async fn dispatch_command( // -- list commands -- Command::LPush { key, values } => { - dispatch!(key, ShardRequest::LPush { key, values }, ResponseTag::LenResultOom) + dispatch!( + key, + ShardRequest::LPush { key, values }, + ResponseTag::LenResultOom + ) } Command::RPush { key, values } => { - dispatch!(key, ShardRequest::RPush { key, values }, ResponseTag::LenResultOom) + dispatch!( + key, + ShardRequest::RPush { key, values }, + ResponseTag::LenResultOom + ) } Command::LPop { key } => { dispatch!(key, ShardRequest::LPop { key }, ResponseTag::PopResult) @@ -790,30 +843,73 @@ async fn dispatch_command( dispatch!(key, ShardRequest::RPop { key }, ResponseTag::PopResult) } Command::LRange { key, start, stop } => { - dispatch!(key, ShardRequest::LRange { key, start, stop }, ResponseTag::ArrayResult) + dispatch!( + key, + ShardRequest::LRange { key, start, stop }, + ResponseTag::ArrayResult + ) } Command::LLen { key } => { dispatch!(key, ShardRequest::LLen { key }, ResponseTag::LenResult) } // -- sorted set commands -- - Command::ZAdd { key, flags, members } => { - dispatch!(key, ShardRequest::ZAdd { - key, members, nx: flags.nx, xx: flags.xx, gt: flags.gt, lt: flags.lt, ch: flags.ch - }, ResponseTag::ZAddResult) + Command::ZAdd { + key, + flags, + members, + } => { + dispatch!( + key, + ShardRequest::ZAdd { + key, + members, + nx: flags.nx, + xx: flags.xx, + gt: flags.gt, + lt: flags.lt, + ch: flags.ch + }, + ResponseTag::ZAddResult + ) } Command::ZRem { key, members } => { - dispatch!(key, ShardRequest::ZRem { key, members }, ResponseTag::ZRemResult) + dispatch!( + key, + ShardRequest::ZRem { key, members }, + ResponseTag::ZRemResult + ) } Command::ZScore { key, member } => { - dispatch!(key, ShardRequest::ZScore { key, member }, ResponseTag::ZScoreResult) + dispatch!( + key, + ShardRequest::ZScore { key, member }, + ResponseTag::ZScoreResult + ) } Command::ZRank { key, member } => { - dispatch!(key, ShardRequest::ZRank { key, member }, ResponseTag::ZRankResult) + dispatch!( + key, + ShardRequest::ZRank { key, member }, + ResponseTag::ZRankResult + ) } - Command::ZRange { key, start, stop, with_scores } => { - dispatch!(key, ShardRequest::ZRange { key, start, stop, with_scores }, - ResponseTag::ZRangeResult { with_scores }) + Command::ZRange { + key, + start, + stop, + with_scores, + } => { + dispatch!( + key, + ShardRequest::ZRange { + key, + start, + stop, + with_scores + }, + ResponseTag::ZRangeResult { with_scores } + ) } Command::ZCard { key } => { dispatch!(key, ShardRequest::ZCard { key }, ResponseTag::LenResult) @@ -821,48 +917,96 @@ async fn dispatch_command( // -- hash commands -- Command::HSet { key, fields } => { - dispatch!(key, ShardRequest::HSet { key, fields }, ResponseTag::HSetResult) + dispatch!( + key, + ShardRequest::HSet { key, fields }, + ResponseTag::HSetResult + ) } Command::HGet { key, field } => { - dispatch!(key, ShardRequest::HGet { key, field }, ResponseTag::HGetResult) + dispatch!( + key, + ShardRequest::HGet { key, field }, + ResponseTag::HGetResult + ) } Command::HGetAll { key } => { - dispatch!(key, ShardRequest::HGetAll { key }, ResponseTag::HGetAllResult) + dispatch!( + key, + ShardRequest::HGetAll { key }, + ResponseTag::HGetAllResult + ) } Command::HDel { key, fields } => { - dispatch!(key, ShardRequest::HDel { key, fields }, ResponseTag::HDelResult) + dispatch!( + key, + ShardRequest::HDel { key, fields }, + ResponseTag::HDelResult + ) } Command::HExists { key, field } => { - dispatch!(key, ShardRequest::HExists { key, field }, ResponseTag::HExistsResult) + dispatch!( + key, + ShardRequest::HExists { key, field }, + ResponseTag::HExistsResult + ) } Command::HLen { key } => { dispatch!(key, ShardRequest::HLen { key }, ResponseTag::LenResult) } Command::HIncrBy { key, field, delta } => { - dispatch!(key, ShardRequest::HIncrBy { key, field, delta }, ResponseTag::HIncrByResult) + dispatch!( + key, + ShardRequest::HIncrBy { key, field, delta }, + ResponseTag::HIncrByResult + ) } Command::HKeys { key } => { - dispatch!(key, ShardRequest::HKeys { key }, ResponseTag::StringArrayResult) + dispatch!( + key, + ShardRequest::HKeys { key }, + ResponseTag::StringArrayResult + ) } Command::HVals { key } => { dispatch!(key, ShardRequest::HVals { key }, ResponseTag::HValsResult) } Command::HMGet { key, fields } => { - dispatch!(key, ShardRequest::HMGet { key, fields }, ResponseTag::HMGetResult) + dispatch!( + key, + ShardRequest::HMGet { key, fields }, + ResponseTag::HMGetResult + ) } // -- set commands -- Command::SAdd { key, members } => { - dispatch!(key, ShardRequest::SAdd { key, members }, ResponseTag::LenResultOom) + dispatch!( + key, + ShardRequest::SAdd { key, members }, + ResponseTag::LenResultOom + ) } Command::SRem { key, members } => { - dispatch!(key, ShardRequest::SRem { key, members }, ResponseTag::LenResult) + dispatch!( + key, + ShardRequest::SRem { key, members }, + ResponseTag::LenResult + ) } Command::SMembers { key } => { - dispatch!(key, ShardRequest::SMembers { key }, ResponseTag::StringArrayResult) + dispatch!( + key, + ShardRequest::SMembers { key }, + ResponseTag::StringArrayResult + ) } Command::SIsMember { key, member } => { - dispatch!(key, ShardRequest::SIsMember { key, member }, ResponseTag::SIsMemberResult) + dispatch!( + key, + ShardRequest::SIsMember { key, member }, + ResponseTag::SIsMemberResult + ) } Command::SCard { key } => { dispatch!(key, ShardRequest::SCard { key }, ResponseTag::LenResult) @@ -870,23 +1014,63 @@ async fn dispatch_command( // -- vector commands -- #[cfg(feature = "vector")] - Command::VAdd { key, element, vector, metric, quantization, connectivity, expansion_add } => { - dispatch!(key, ShardRequest::VAdd { - key, element, vector, metric, quantization, connectivity, expansion_add - }, ResponseTag::VAddResult) + Command::VAdd { + key, + element, + vector, + metric, + quantization, + connectivity, + expansion_add, + } => { + dispatch!( + key, + ShardRequest::VAdd { + key, + element, + vector, + metric, + quantization, + connectivity, + expansion_add + }, + ResponseTag::VAddResult + ) } #[cfg(feature = "vector")] - Command::VSim { key, query, count, ef_search, with_scores } => { - dispatch!(key, ShardRequest::VSim { key, query, count, ef_search }, - ResponseTag::VSimResult { with_scores }) + Command::VSim { + key, + query, + count, + ef_search, + with_scores, + } => { + dispatch!( + key, + ShardRequest::VSim { + key, + query, + count, + ef_search + }, + ResponseTag::VSimResult { with_scores } + ) } #[cfg(feature = "vector")] Command::VRem { key, element } => { - dispatch!(key, ShardRequest::VRem { key, element }, ResponseTag::VRemResult) + dispatch!( + key, + ShardRequest::VRem { key, element }, + ResponseTag::VRemResult + ) } #[cfg(feature = "vector")] Command::VGet { key, element } => { - dispatch!(key, ShardRequest::VGet { key, element }, ResponseTag::VGetResult) + dispatch!( + key, + ShardRequest::VGet { key, element }, + ResponseTag::VGetResult + ) } #[cfg(feature = "vector")] Command::VCard { key } => { @@ -908,25 +1092,38 @@ async fn dispatch_command( "ERR source and destination keys must hash to the same shard".into(), )) } else { - dispatch!(key, ShardRequest::Rename { key, newkey }, ResponseTag::RenameResult) + dispatch!( + key, + ShardRequest::Rename { key, newkey }, + ResponseTag::RenameResult + ) } } // -- proto commands that are single-key dispatches -- #[cfg(feature = "protobuf")] - Command::ProtoSet { key, type_name, data, expire, nx, xx } => { + Command::ProtoSet { + key, + type_name, + data, + expire, + nx, + xx, + } => { if engine.schema_registry().is_none() { - return PendingResponse::Immediate( - Frame::Error("ERR protobuf support is not enabled".into()), - ); + return PendingResponse::Immediate(Frame::Error( + "ERR protobuf support is not enabled".into(), + )); } let registry = engine.schema_registry().unwrap(); { let reg = match registry.read() { Ok(r) => r, - Err(_) => return PendingResponse::Immediate( - Frame::Error("ERR schema registry lock poisoned".into()), - ), + Err(_) => { + return PendingResponse::Immediate(Frame::Error( + "ERR schema registry lock poisoned".into(), + )) + } }; if let Err(e) = reg.validate(&type_name, &data) { return PendingResponse::Immediate(Frame::Error(format!("ERR {e}"))); @@ -936,47 +1133,78 @@ async fn dispatch_command( SetExpire::Ex(secs) => Duration::from_secs(secs), SetExpire::Px(millis) => Duration::from_millis(millis), }); - dispatch!(key, ShardRequest::ProtoSet { - key, type_name, data, expire: duration, nx, xx - }, ResponseTag::ProtoSetResult) + dispatch!( + key, + ShardRequest::ProtoSet { + key, + type_name, + data, + expire: duration, + nx, + xx + }, + ResponseTag::ProtoSetResult + ) } #[cfg(feature = "protobuf")] Command::ProtoGet { key } => { if engine.schema_registry().is_none() { - return PendingResponse::Immediate( - Frame::Error("ERR protobuf support is not enabled".into()), - ); + return PendingResponse::Immediate(Frame::Error( + "ERR protobuf support is not enabled".into(), + )); } - dispatch!(key, ShardRequest::ProtoGet { key }, ResponseTag::ProtoGetResult) + dispatch!( + key, + ShardRequest::ProtoGet { key }, + ResponseTag::ProtoGetResult + ) } #[cfg(feature = "protobuf")] Command::ProtoType { key } => { if engine.schema_registry().is_none() { - return PendingResponse::Immediate( - Frame::Error("ERR protobuf support is not enabled".into()), - ); + return PendingResponse::Immediate(Frame::Error( + "ERR protobuf support is not enabled".into(), + )); } - dispatch!(key, ShardRequest::ProtoType { key }, ResponseTag::ProtoTypeResult) + dispatch!( + key, + ShardRequest::ProtoType { key }, + ResponseTag::ProtoTypeResult + ) } #[cfg(feature = "protobuf")] - Command::ProtoSetField { key, field_path, value } => { + Command::ProtoSetField { + key, + field_path, + value, + } => { if engine.schema_registry().is_none() { - return PendingResponse::Immediate( - Frame::Error("ERR protobuf support is not enabled".into()), - ); + return PendingResponse::Immediate(Frame::Error( + "ERR protobuf support is not enabled".into(), + )); } - dispatch!(key, ShardRequest::ProtoSetField { key, field_path, value }, - ResponseTag::ProtoSetFieldResult) + dispatch!( + key, + ShardRequest::ProtoSetField { + key, + field_path, + value + }, + ResponseTag::ProtoSetFieldResult + ) } #[cfg(feature = "protobuf")] Command::ProtoDelField { key, field_path } => { if engine.schema_registry().is_none() { - return PendingResponse::Immediate( - Frame::Error("ERR protobuf support is not enabled".into()), - ); + return PendingResponse::Immediate(Frame::Error( + "ERR protobuf support is not enabled".into(), + )); } - dispatch!(key, ShardRequest::ProtoDelField { key, field_path }, - ResponseTag::ProtoDelFieldResult) + dispatch!( + key, + ShardRequest::ProtoDelField { key, field_path }, + ResponseTag::ProtoDelFieldResult + ) } // -- everything else falls back to the full execute() path -- @@ -995,7 +1223,12 @@ async fn resolve_response( ) -> Frame { match pending { PendingResponse::Immediate(frame) => frame, - PendingResponse::Pending { rx, tag, start, cmd_name } => { + PendingResponse::Pending { + rx, + tag, + start, + cmd_name, + } => { let frame = match rx.await { Ok(resp) => resolve_shard_response(resp, tag), Err(_) => Frame::Error("ERR shard unavailable".into()), @@ -1175,9 +1408,7 @@ fn resolve_shard_response(resp: ShardResponse, tag: ResponseTag) -> Frame { other => Frame::Error(format!("ERR unexpected shard response: {other:?}")), }, ResponseTag::HValsResult => match resp { - ShardResponse::Array(vals) => { - Frame::Array(vals.into_iter().map(Frame::Bulk).collect()) - } + ShardResponse::Array(vals) => Frame::Array(vals.into_iter().map(Frame::Bulk).collect()), ShardResponse::WrongType => wrongtype_error(), other => Frame::Error(format!("ERR unexpected shard response: {other:?}")), }, @@ -1205,9 +1436,7 @@ fn resolve_shard_response(resp: ShardResponse, tag: ResponseTag) -> Frame { }, #[cfg(feature = "vector")] ResponseTag::VAddResult => match resp { - ShardResponse::VAddResult { added, .. } => { - Frame::Integer(if added { 1 } else { 0 }) - } + ShardResponse::VAddResult { added, .. } => Frame::Integer(if added { 1 } else { 0 }), ShardResponse::WrongType => wrongtype_error(), ShardResponse::OutOfMemory => oom_error(), ShardResponse::Err(msg) => Frame::Error(format!("ERR {msg}")), diff --git a/tests/integration/src/helpers.rs b/tests/integration/src/helpers.rs index b0b1581f..79874579 100644 --- a/tests/integration/src/helpers.rs +++ b/tests/integration/src/helpers.rs @@ -56,7 +56,8 @@ impl TestServer { let mut cmd = Command::new(&binary); cmd.arg("--port").arg(port.to_string()); cmd.arg("--host").arg("127.0.0.1"); - cmd.arg("--shards").arg(opts.shards.unwrap_or(2).to_string()); + cmd.arg("--shards") + .arg(opts.shards.unwrap_or(2).to_string()); // suppress tracing output in tests cmd.env("RUST_LOG", "error"); @@ -160,6 +161,7 @@ impl TestClient { /// Sends a command with raw byte arguments and returns the parsed response. /// Useful for binary data like protobuf descriptors. + #[allow(dead_code)] pub async fn cmd_raw(&mut self, args: &[&[u8]]) -> Frame { let parts: Vec = args .iter()