diff --git a/README.md b/README.md index 52f358eee..007c0b107 100644 --- a/README.md +++ b/README.md @@ -51,7 +51,9 @@ Moon implements 200+ Redis commands with a thread-per-core shared-nothing archit Shard Scaling & Production Value

-Benchmarked against Redis 8.6.1 on Apple M4 Pro (co-located, `redis-benchmark`): +> **Note:** All benchmarks are run on **ARM64** hardware — Apple M4 Pro (macOS) and OrbStack Linux VM (aarch64). x86_64 results may differ due to architectural differences in io_uring, memory subsystem, and SIMD paths. + +#### macOS ARM64 (Apple M4 Pro, co-located, Redis 8.6.1) | Metric | Moon vs Redis | Conditions | |--------|:------------:|------------| @@ -63,7 +65,19 @@ Benchmarked against Redis 8.6.1 on Apple M4 Pro (co-located, `redis-benchmark`): | Memory (1KB+ values) | **27-35% less** | Per-key RSS measurement | | p50 latency (8 shards) | **8-10x lower** | 0.031ms vs 0.26ms | | CPU efficiency (p=64) | **45x better** | 1.9% vs 43.9% CPU | -| Data correctness | **132/132 tests** | All types, 1/4/12 shards | + +#### Linux ARM64 (OrbStack aarch64, Ubuntu 25.10, Redis 8.0.2, monoio io_uring) + +| Metric | Moon vs Redis | Conditions | +|--------|:------------:|------------| +| Peak GET throughput | **11.7M ops/sec** | 1 shard, pipeline=64 | +| Peak SET throughput | **5.4M ops/sec** | 1 shard, pipeline=64 | +| SET with AOF (p=64) | **2.15x faster** | Per-shard WAL, appendfsync=everysec | +| GET with AOF (p=64) | **2.17x faster** | io_uring batched submission | +| SET with AOF (p=16) | **2.06x faster** | Per-shard WAL eliminates global bottleneck | +| Crash recovery | **100%** | 38/38 consistency checks, all data types | +| RDB load speed | **5.9M keys/sec** | 632K keys in 107ms | +| AOF compaction | **1.16x Redis** | Multi-part AOF (base.rdb + incr.aof) | See [BENCHMARK.md](BENCHMARK.md) for full methodology and results, or [BENCHMARK-PRODUCTION.md](BENCHMARK-PRODUCTION.md) for production workload patterns. @@ -85,9 +99,11 @@ See [BENCHMARK.md](BENCHMARK.md) for full methodology and results, or [BENCHMARK - **Lock-free channels** - Custom oneshot channels replacing tokio::oneshot (12% CPU reduction) ### Persistence +- **Multi-part AOF** - Inspired by Redis 7+ design: `moon.aof..base.rdb` + `moon.aof..incr.aof` + `moon.aof.manifest` in `appendonlydir/`. Note: Moon's AOF format is **not compatible** with Redis — file naming, manifest format, and RDB encoding differ. Direct migration between Moon and Redis AOF files is not supported. +- **BGREWRITEAOF** - RDB preamble compaction, automatic old file cleanup, 100% crash recovery - **RDB snapshots** - Forkless compartmentalized snapshots (no COW memory spike) -- **AOF** - Per-shard WAL with batched fsync, configurable everysec/always/no -- **WAL v2** - Checksums, block framing, corruption isolation +- **Per-shard WAL** - CRC32-checksummed block frames, configurable everysec/always/no fsync +- **Fast RDB loader** - Significantly faster bulk loading than Redis (pre-sized hash tables, direct deserialization, skipped duplicate checks during restore) ### Networking & Protocol - **RESP2/RESP3** - Full protocol support with HELLO negotiation diff --git a/scripts/test-consistency.sh b/scripts/test-consistency.sh deleted file mode 100755 index f854cb508..000000000 --- a/scripts/test-consistency.sh +++ /dev/null @@ -1,523 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -############################################################################### -# test-consistency.sh -- Data consistency test: SET/GET, SETEX/GETEX, collections -# -# Verifies that data written to moon can be read back identically. -# Tests all size ranges (SSO inline, heap small, heap large, binary). -# Compares moon output against Redis as ground truth. -# -# Usage: -# ./scripts/test-consistency.sh [--shards N] [--skip-build] [--port-rust N] -############################################################################### - -PORT_REDIS=6399 -PORT_RUST=6400 -SHARDS=1 -SKIP_BUILD=false -RUST_BINARY="./target/release/moon" -PASS=0 -FAIL=0 -RUST_PID="" -REDIS_PID="" - -while [[ $# -gt 0 ]]; do - case "$1" in - --shards) SHARDS="$2"; shift 2 ;; - --skip-build) SKIP_BUILD=true; shift ;; - --port-rust) PORT_RUST="$2"; shift 2 ;; - *) echo "Unknown: $1"; exit 1 ;; - esac -done - -log() { echo "[$(date '+%H:%M:%S')] $*"; } - -cleanup() { - [[ -n "${RUST_PID:-}" ]] && kill "$RUST_PID" 2>/dev/null; wait "$RUST_PID" 2>/dev/null || true - [[ -n "${REDIS_PID:-}" ]] && kill "$REDIS_PID" 2>/dev/null; wait "$REDIS_PID" 2>/dev/null || true - pkill -f "redis-server.*${PORT_REDIS}" 2>/dev/null || true - pkill -f "moon.*${PORT_RUST}" 2>/dev/null || true -} -trap cleanup EXIT - -assert_eq() { - local desc="$1" expected="$2" actual="$3" - if [[ "$expected" == "$actual" ]]; then - PASS=$((PASS + 1)) - else - FAIL=$((FAIL + 1)) - echo " FAIL: $desc" - echo " expected: $(echo "$expected" | head -c 200)" - echo " actual: $(echo "$actual" | head -c 200)" - fi -} - -# Run same command on both servers, compare output -assert_both() { - local desc="$1"; shift - local redis_out rust_out - redis_out=$(redis-cli -p "$PORT_REDIS" "$@" 2>&1) || true - rust_out=$(redis-cli -p "$PORT_RUST" "$@" 2>&1) || true - assert_eq "$desc" "$redis_out" "$rust_out" -} - -# Run commands on both servers (no comparison, just execute) -both() { - redis-cli -p "$PORT_REDIS" "$@" &>/dev/null || true - redis-cli -p "$PORT_RUST" "$@" &>/dev/null || true -} - -wait_for_port() { - local port=$1 - for ((i=0; i<30; i++)); do - redis-cli -p "$port" PING 2>/dev/null | grep -q PONG && return 0 - sleep 0.2 - done - log "ERROR: port $port not ready"; return 1 -} - -# =========================================================================== -# Setup -# =========================================================================== - -if [[ "$SKIP_BUILD" == false ]]; then - log "Building..." - RUSTFLAGS="-C target-cpu=native" cargo build --release 2>&1 | tail -2 -fi - -log "Starting Redis on :$PORT_REDIS ..." -redis-server --port "$PORT_REDIS" --save "" --appendonly no --loglevel warning --daemonize no &>/dev/null & -REDIS_PID=$! - -log "Starting moon on :$PORT_RUST (shards=$SHARDS)..." -"$RUST_BINARY" --port "$PORT_RUST" --shards "$SHARDS" &>/dev/null & -RUST_PID=$! - -wait_for_port "$PORT_REDIS" -wait_for_port "$PORT_RUST" -both FLUSHALL - -# =========================================================================== -# 1. String SET/GET — size ranges -# =========================================================================== -log "=== 1. String SET/GET size ranges ===" - -# Empty string -both SET str:empty "" -assert_both "GET empty string" GET str:empty - -# 1 byte -both SET str:1b "x" -assert_both "GET 1-byte" GET str:1b - -# 12 bytes (max SSO inline) -both SET str:12b "123456789012" -assert_both "GET 12-byte (SSO boundary)" GET str:12b - -# 13 bytes (first heap path) -both SET str:13b "1234567890123" -assert_both "GET 13-byte (heap boundary)" GET str:13b - -# 64 bytes -VAL64=$(python3 -c "print('A' * 64)") -both SET str:64b "$VAL64" -assert_both "GET 64-byte" GET str:64b - -# 256 bytes -VAL256=$(python3 -c "print('B' * 256)") -both SET str:256b "$VAL256" -assert_both "GET 256-byte" GET str:256b - -# 1KB -VAL1K=$(python3 -c "print('C' * 1024)") -both SET str:1k "$VAL1K" -assert_both "GET 1KB" GET str:1k - -# 4KB -VAL4K=$(python3 -c "print('D' * 4096)") -both SET str:4k "$VAL4K" -assert_both "GET 4KB" GET str:4k - -# 64KB -VAL64K=$(python3 -c "print('E' * 65536)") -both SET str:64k "$VAL64K" -assert_both "GET 64KB" GET str:64k - -# Numeric string -both SET str:num "1234567890" -assert_both "GET numeric string" GET str:num - -# Negative number -both SET str:neg "-99999" -assert_both "GET negative number" GET str:neg - -# Float -both SET str:float "3.14159265358979" -assert_both "GET float string" GET str:float - -# =========================================================================== -# 2. String mutations -# =========================================================================== -log "=== 2. String mutations ===" - -# APPEND -both SET mut:append "hello" -both APPEND mut:append " world" -assert_both "APPEND result" GET mut:append - -# APPEND crossing SSO boundary (start <12, end >12) -both SET mut:cross "12345678901" # 11 bytes (SSO) -both APPEND mut:cross "XY" # 13 bytes (heap) -assert_both "APPEND SSO->heap" GET mut:cross - -# INCR / DECR -both SET mut:counter "100" -both INCR mut:counter -assert_both "INCR" GET mut:counter -both DECR mut:counter -both DECR mut:counter -assert_both "DECR twice" GET mut:counter -both INCRBY mut:counter 50 -assert_both "INCRBY 50" GET mut:counter - -# INCRBYFLOAT (skip exact comparison — float formatting may differ) -both SET mut:flt "10.5" -both INCRBYFLOAT mut:flt "0.1" -rust_flt=$(redis-cli -p "$PORT_RUST" GET mut:flt 2>&1) -if [[ "$rust_flt" == "10.6" || "$rust_flt" == "10.59999999999999964" ]]; then - PASS=$((PASS + 1)) -else - FAIL=$((FAIL + 1)); echo " FAIL: INCRBYFLOAT unexpected: $rust_flt" -fi - -# GETRANGE (may not be implemented — test only if supported) -both SET mut:range "Hello, World!" -rust_gr=$(redis-cli -p "$PORT_RUST" GETRANGE mut:range 0 4 2>&1) -if [[ "$rust_gr" != *"unknown command"* ]]; then - assert_both "GETRANGE 0 4" GETRANGE mut:range 0 4 - assert_both "GETRANGE 7 -1" GETRANGE mut:range 7 -1 -else - log " SKIP: GETRANGE not implemented" -fi - -# SETRANGE (may not be implemented) -both SET mut:setrange "Hello, World!" -rust_sr=$(redis-cli -p "$PORT_RUST" SETRANGE mut:setrange 7 "Redis" 2>&1) -if [[ "$rust_sr" != *"unknown command"* ]]; then - assert_both "SETRANGE" GET mut:setrange -else - log " SKIP: SETRANGE not implemented" -fi - -# STRLEN -assert_both "STRLEN 13-byte" STRLEN str:13b -assert_both "STRLEN 1KB" STRLEN str:1k - -# GETDEL -both SET mut:getdel "deleteme" -assert_both "GETDEL returns value" GETDEL mut:getdel -assert_both "GETDEL key gone" GET mut:getdel - -# GETSET (deprecated but still valid) -both SET mut:getset "old" -assert_both "GETSET returns old" GETSET mut:getset "new" -assert_both "GETSET new value" GET mut:getset - -# =========================================================================== -# 3. MSET / MGET -# =========================================================================== -log "=== 3. MSET / MGET ===" - -both MSET mk1 "val1" mk2 "val2" mk3 "val3" -assert_both "MGET 3 keys" MGET mk1 mk2 mk3 -assert_both "MGET with missing" MGET mk1 nonexistent mk3 - -# =========================================================================== -# 4. SET with options (EX, PX, NX, XX, KEEPTTL, GET) -# =========================================================================== -log "=== 4. SET with options ===" - -both SET opt:ex "expire_me" EX 3600 -assert_both "SET EX value" GET opt:ex -# TTL should be close to 3600 -redis_ttl=$(redis-cli -p "$PORT_REDIS" TTL opt:ex) -rust_ttl=$(redis-cli -p "$PORT_RUST" TTL opt:ex) -if (( rust_ttl >= 3598 && rust_ttl <= 3600 )); then - PASS=$((PASS + 1)) -else - FAIL=$((FAIL + 1)) - echo " FAIL: TTL mismatch: redis=$redis_ttl rust=$rust_ttl" -fi - -both SET opt:px "px_value" PX 60000 -assert_both "SET PX value" GET opt:px - -# NX (set only if not exists) -both SET opt:nx "original" -both SET opt:nx "overwrite" NX # should fail -assert_both "SET NX no overwrite" GET opt:nx - -# XX (set only if exists) -both SET opt:xx "impossible" XX # key doesn't exist, should fail for new key -both SET opt:xxreal "first" -both SET opt:xxreal "second" XX # should succeed -assert_both "SET XX overwrites" GET opt:xxreal - -# SETEX / SETNX -both SETEX setex:key 3600 "setex_value" -assert_both "SETEX value" GET setex:key - -both DEL setnx:key -both SETNX setnx:key "first" -both SETNX setnx:key "second" # should fail -assert_both "SETNX no overwrite" GET setnx:key - -# =========================================================================== -# 5. Binary-safe data -# =========================================================================== -log "=== 5. Binary-safe data ===" - -# Use redis-cli with hex to set binary values -redis-cli -p "$PORT_REDIS" SET bin:null $'\x00\x01\x02\x03' &>/dev/null || true -redis-cli -p "$PORT_RUST" SET bin:null $'\x00\x01\x02\x03' &>/dev/null || true -assert_both "Binary with null bytes" GET bin:null - -# Special characters -both SET bin:special "hello\tworld\nnewline" -assert_both "Tab and newline" GET bin:special - -both SET bin:utf8 "Hello" -assert_both "UTF-8 emoji" GET bin:utf8 - -# =========================================================================== -# 6. Hash SET/GET -# =========================================================================== -log "=== 6. Hash operations ===" - -both HSET h:test f1 "val1" f2 "val2" f3 "val3" -assert_both "HGET f1" HGET h:test f1 -assert_both "HGET f2" HGET h:test f2 -# HGETALL order may differ — sort for comparison -redis_hga=$(redis-cli -p "$PORT_REDIS" HGETALL h:test 2>&1 | sort) -rust_hga=$(redis-cli -p "$PORT_RUST" HGETALL h:test 2>&1 | sort) -assert_eq "HGETALL (sorted)" "$redis_hga" "$rust_hga" -assert_both "HMGET" HMGET h:test f1 f3 nonexistent -assert_both "HLEN" HLEN h:test -assert_both "HEXISTS f1" HEXISTS h:test f1 -assert_both "HEXISTS missing" HEXISTS h:test missing - -# Large hash value -HVAL=$(python3 -c "print('X' * 1024)") -both HSET h:test f_large "$HVAL" -assert_both "HGET large value" HGET h:test f_large - -both HDEL h:test f2 -assert_both "HDEL then HGET" HGET h:test f2 -assert_both "HLEN after HDEL" HLEN h:test - -both HINCRBY h:test counter 10 -assert_both "HINCRBY" HGET h:test counter -both HINCRBY h:test counter 5 -assert_both "HINCRBY again" HGET h:test counter - -# =========================================================================== -# 7. List operations -# =========================================================================== -log "=== 7. List operations ===" - -both RPUSH l:test a b c d e -assert_both "LRANGE all" LRANGE l:test 0 -1 -assert_both "LLEN" LLEN l:test -assert_both "LINDEX 0" LINDEX l:test 0 -assert_both "LINDEX -1" LINDEX l:test -1 - -both LPUSH l:test z -assert_both "LPUSH + LRANGE" LRANGE l:test 0 -1 - -both RPOP l:test -assert_both "RPOP + LRANGE" LRANGE l:test 0 -1 - -both LPOP l:test -assert_both "LPOP + LRANGE" LRANGE l:test 0 -1 - -# Large list values -LVAL=$(python3 -c "print('Y' * 512)") -both RPUSH l:test "$LVAL" -assert_both "LINDEX large value" LINDEX l:test -1 - -# =========================================================================== -# 8. Set operations -# =========================================================================== -log "=== 8. Set operations ===" - -both SADD s:test a b c d e -assert_both "SCARD" SCARD s:test -assert_both "SISMEMBER a" SISMEMBER s:test a -assert_both "SISMEMBER missing" SISMEMBER s:test z - -both SREM s:test c -assert_both "SCARD after SREM" SCARD s:test -assert_both "SISMEMBER removed" SISMEMBER s:test c - -# SMEMBERS order may differ — sort both -redis_sm=$(redis-cli -p "$PORT_REDIS" SMEMBERS s:test 2>&1 | sort) -rust_sm=$(redis-cli -p "$PORT_RUST" SMEMBERS s:test 2>&1 | sort) -assert_eq "SMEMBERS (sorted)" "$redis_sm" "$rust_sm" - -# =========================================================================== -# 9. Sorted Set operations -# =========================================================================== -log "=== 9. Sorted Set operations ===" - -both ZADD z:test 1.0 "alpha" 2.5 "beta" 3.0 "gamma" 0.5 "delta" -assert_both "ZCARD" ZCARD z:test -assert_both "ZSCORE alpha" ZSCORE z:test alpha -assert_both "ZSCORE beta" ZSCORE z:test beta -assert_both "ZRANK alpha" ZRANK z:test alpha -assert_both "ZRANGE 0 -1" ZRANGE z:test 0 -1 -assert_both "ZRANGE WITHSCORES" ZRANGE z:test 0 -1 WITHSCORES -assert_both "ZRANGEBYSCORE 1 3" ZRANGEBYSCORE z:test 1 3 - -both ZINCRBY z:test 10 "delta" -assert_both "ZINCRBY then ZSCORE" ZSCORE z:test delta - -# =========================================================================== -# 10. Bulk data consistency (redis-benchmark load + random verify) -# =========================================================================== -log "=== 10. Bulk data consistency (1K deterministic keys) ===" - -both FLUSHALL - -# Deterministic load: 1K keys with varied value sizes -for i in $(seq 0 999); do - key="bulk:$(printf '%04d' "$i")" - # Vary sizes: 0-255 bytes padding - pad=$(python3 -c "print('x' * ($i % 256))") - val="v${i}_${pad}" - both SET "$key" "$val" -done - -# DBSIZE: verify both have 1000 keys (exact match not required due to prior test keys) -redis_db=$(redis-cli -p "$PORT_REDIS" DBSIZE 2>&1 | grep -oE '[0-9]+') -rust_db=$(redis-cli -p "$PORT_RUST" DBSIZE 2>&1 | grep -oE '[0-9]+') -if (( redis_db >= 1000 && rust_db >= 1000 )); then - PASS=$((PASS + 1)) -else - FAIL=$((FAIL + 1)); echo " FAIL: DBSIZE: redis=$redis_db rust=$rust_db (expected >=1000)" -fi - -# Spot-check 50 random keys -BULK_PASS=0 -BULK_FAIL=0 -for i in $(python3 -c "import random; random.seed(42); print(' '.join(str(random.randint(0,999)) for _ in range(50)))"); do - key="bulk:$(printf '%04d' "$i")" - rv=$(redis-cli -p "$PORT_REDIS" GET "$key" 2>&1) - uv=$(redis-cli -p "$PORT_RUST" GET "$key" 2>&1) - if [[ "$rv" == "$uv" ]]; then - BULK_PASS=$((BULK_PASS + 1)) - else - BULK_FAIL=$((BULK_FAIL + 1)) - echo " FAIL: bulk $key" - echo " redis: $(echo "$rv" | head -c 100)" - echo " rust: $(echo "$uv" | head -c 100)" - fi -done -PASS=$((PASS + BULK_PASS)) -FAIL=$((FAIL + BULK_FAIL)) -log " Bulk spot-check: $BULK_PASS/$((BULK_PASS + BULK_FAIL)) passed" - -# =========================================================================== -# 11. Overwrite consistency -# =========================================================================== -log "=== 11. Overwrite / type change ===" - -# Overwrite string with different sizes -both SET ow:key "small" -assert_both "GET before overwrite" GET ow:key -both SET ow:key "$VAL1K" -assert_both "GET after overwrite with 1KB" GET ow:key -both SET ow:key "tiny" -assert_both "GET after shrink overwrite" GET ow:key - -# Overwrite with different type -both DEL ow:type -both SET ow:type "string_val" -assert_both "GET string" GET ow:type -both DEL ow:type -both HSET ow:type f1 v1 -assert_both "HGET after type change" HGET ow:type f1 - -# =========================================================================== -# 12. Edge cases -# =========================================================================== -log "=== 12. Edge cases ===" - -# GET nonexistent key -assert_both "GET nonexistent" GET totally:missing:key - -# DEL + GET -both SET edge:del "exists" -both DEL edge:del -assert_both "GET after DEL" GET edge:del - -# SETNX on existing -both SET edge:setnx "original" -both SET edge:setnx "new" NX -assert_both "SETNX on existing" GET edge:setnx - -# SET with GET option -both SET edge:setget "old_value" -assert_both "SET GET returns old" SET edge:setget "new_value" GET -assert_both "SET GET new value" GET edge:setget - -# Very long key name -LONGKEY=$(python3 -c "print('k' * 500)") -both SET "$LONGKEY" "long_key_value" -assert_both "GET with 500-char key" GET "$LONGKEY" - -# =========================================================================== -# Summary -# =========================================================================== - -echo "" -# =========================================================================== -# Vector Search (moon-only — FT.* not available in Redis) -# =========================================================================== -log "=== Vector Search (moon-only) ===" - -# Create index on moon only -FT_CREATE=$(redis-cli -p "$PORT_RUST" FT.CREATE vecidx ON HASH PREFIX 1 vec: SCHEMA embedding VECTOR FLAT 6 DIM 4 DISTANCE_METRIC L2 TYPE FLOAT32 2>&1) -assert_eq "FT.CREATE" "OK" "$FT_CREATE" - -# Insert vectors — use python3 to avoid null byte stripping in bash -python3 -c "import struct,sys; sys.stdout.buffer.write(struct.pack('<4f',1.0,0.0,0.0,0.0))" | redis-cli -x -p "$PORT_RUST" HSET vec:1 embedding >/dev/null 2>&1 -python3 -c "import struct,sys; sys.stdout.buffer.write(struct.pack('<4f',0.0,1.0,0.0,0.0))" | redis-cli -x -p "$PORT_RUST" HSET vec:2 embedding >/dev/null 2>&1 - -# FT.INFO should show index -FT_INFO=$(redis-cli -p "$PORT_RUST" FT.INFO vecidx 2>&1) -if echo "$FT_INFO" | grep -q "vecidx"; then - PASS=$((PASS + 1)) -else - FAIL=$((FAIL + 1)); echo " FAIL: FT.INFO should show vecidx" -fi - -# FT.DROPINDEX -FT_DROP=$(redis-cli -p "$PORT_RUST" FT.DROPINDEX vecidx 2>&1) -assert_eq "FT.DROPINDEX" "OK" "$FT_DROP" - -echo "============================================" -echo " Data Consistency Test Results" -echo "============================================" -echo " PASSED: $PASS" -echo " FAILED: $FAIL" -echo " TOTAL: $((PASS + FAIL))" -echo "============================================" - -if (( FAIL > 0 )); then - echo " STATUS: FAIL" - exit 1 -else - echo " STATUS: ALL PASSED" - exit 0 -fi diff --git a/src/command/persistence.rs b/src/command/persistence.rs index 928a2b1fe..35423fabf 100644 --- a/src/command/persistence.rs +++ b/src/command/persistence.rs @@ -212,6 +212,21 @@ pub fn bgrewriteaof_start(aof_tx: &channel::MpscSender, db: SharedDa } } +/// Start BGREWRITEAOF in sharded mode using ShardDatabases. +pub fn bgrewriteaof_start_sharded( + aof_tx: &channel::MpscSender, + shard_databases: std::sync::Arc, +) -> Frame { + match aof_tx.try_send(AofMessage::RewriteSharded(shard_databases)) { + Ok(()) => Frame::SimpleString(Bytes::from_static( + b"Background append only file rewriting started", + )), + Err(_) => Frame::Error(Bytes::from_static( + b"ERR Background AOF rewrite failed to start", + )), + } +} + /// SAVE command: synchronous save to disk. Blocks until complete. /// /// Clones all entries under read locks (same as BGSAVE), then serializes diff --git a/src/main.rs b/src/main.rs index f45ba3182..54d605f15 100644 --- a/src/main.rs +++ b/src/main.rs @@ -203,17 +203,74 @@ fn main() -> anyhow::Result<()> { // Create and restore all shards on main thread, then extract databases // into centralized ShardDatabases for cross-shard direct read access. + let skip_shard_wal = config.appendonly == "yes"; let mut shards: Vec = (0..num_shards) .map(|id| { let mut shard = Shard::new(id, num_shards, config.databases, config.to_runtime_config()); if let Some(ref dir) = persistence_dir { - shard.restore_from_persistence(dir); + shard.restore_from_persistence(dir, skip_shard_wal); } shard }) .collect(); + // Replay AOF data from disk (supplements per-shard WAL restore above). + // + // Priority order: + // 1. Multi-part AOF (appendonlydir/ with manifest) — base RDB + incremental RESP + // 2. Legacy single-file AOF (appendonly.aof) — RDB preamble or pure RESP + if config.appendonly == "yes" { + if let Some(ref dir) = persistence_dir { + use moon::persistence::aof_manifest::AofManifest; + use moon::persistence::replay::DispatchReplayEngine; + + let base_dir = std::path::PathBuf::from(dir); + let target_dbs = &mut shards[0].databases; + + if let Some(manifest) = AofManifest::load(&base_dir) { + // Multi-part AOF: load base RDB + replay incremental RESP + match moon::persistence::aof_manifest::replay_multi_part( + target_dbs, + &manifest, + &DispatchReplayEngine, + ) { + Ok(n) => info!( + "AOF loaded (multi-part seq {}): {} keys/commands", + manifest.seq, n + ), + Err(e) => tracing::error!("AOF multi-part load failed: {}", e), + } + } else { + // Legacy single-file AOF (backward compatible) + let aof_path = base_dir.join(&config.appendfilename); + if aof_path.exists() { + match aof::replay_aof(target_dbs, &aof_path, &DispatchReplayEngine) { + Ok(n) => info!( + "AOF loaded (legacy): {} commands from {}", + n, + aof_path.display() + ), + Err(e) => tracing::error!("AOF load failed: {}", e), + } + } + } + } + + // Ensure multi-part AOF manifest exists for the writer thread. + // Recovery is now complete, so it's safe to create the manifest + // without racing against legacy AOF migration detection. + if let Some(ref dir) = persistence_dir { + use moon::persistence::aof_manifest::AofManifest; + let base_dir = std::path::PathBuf::from(dir); + if AofManifest::load(&base_dir).is_none() { + if let Err(e) = AofManifest::initialize(&base_dir) { + tracing::error!("Failed to initialize AOF manifest after recovery: {}", e); + } + } + } + } + // Extract databases from all shards and wrap in ShardDatabases let all_dbs: Vec> = shards .iter_mut() diff --git a/src/persistence/aof.rs b/src/persistence/aof.rs index b0161ca78..5caeb4d1e 100644 --- a/src/persistence/aof.rs +++ b/src/persistence/aof.rs @@ -61,6 +61,8 @@ pub enum AofMessage { Append(Bytes), /// Trigger a full AOF rewrite (compaction) using current database state. Rewrite(SharedDatabases), + /// Trigger AOF rewrite in sharded mode (all shards' databases). + RewriteSharded(Arc), /// Shut down the AOF writer task gracefully. Shutdown, } @@ -107,38 +109,129 @@ pub async fn aof_writer_task( #[cfg(feature = "runtime-tokio")] interval.tick().await; // consume first tick - // Monoio fallback: AOF writer uses sync I/O in a simple recv loop. + // Monoio path: multi-part AOF (base RDB + incremental RESP) with sync I/O. + // + // On startup, if appendonlydir/ exists with a manifest, open the current + // incr file for appending. Otherwise start fresh with seq 1. + // On BGREWRITEAOF: snapshot → write new base RDB → create new incr → advance manifest. #[cfg(feature = "runtime-monoio")] { + use crate::persistence::aof_manifest::AofManifest; use std::io::Write; + + // Resolve the persistence base directory from aof_path's parent. + let base_dir = aof_path.parent().unwrap_or(Path::new(".")).to_path_buf(); + + // Load manifest — do NOT create one here if it doesn't exist. + // main.rs recovery runs concurrently and must finish before a manifest + // is created, to avoid racing against legacy single-file AOF detection. + // main.rs will create the manifest after recovery completes. + let mut manifest = loop { + if let Some(m) = AofManifest::load(&base_dir) { + break m; + } + // main.rs recovery hasn't created the manifest yet — wait. + std::thread::sleep(std::time::Duration::from_millis(50)); + }; + + // Open the current incremental file for appending + let incr_path = manifest.incr_path(); let mut file = match std::fs::OpenOptions::new() .create(true) .append(true) - .open(&aof_path) + .open(&incr_path) { Ok(f) => f, Err(e) => { - error!("Failed to open AOF file {}: {}", aof_path.display(), e); + error!( + "Failed to open AOF incr file {}: {}", + incr_path.display(), + e + ); return; } }; + info!( + "AOF writer: seq {}, incr={}", + manifest.seq, + incr_path.display() + ); + + let mut last_fsync = Instant::now(); + + let mut write_error = false; + loop { - // Use blocking recv since monoio doesn't support tokio::select! match rx.recv() { Ok(AofMessage::Append(data)) => { - let _ = file.write_all(&data); - if fsync == FsyncPolicy::Always { - let _ = file.flush(); - let _ = std::io::Write::flush(&mut file); + if write_error { + continue; // Drop appends after persistent I/O failure + } + if let Err(e) = file.write_all(&data) { + error!( + "AOF write failed (seq {}): {}. Persistence degraded.", + manifest.seq, e + ); + write_error = true; + continue; + } + match fsync { + FsyncPolicy::Always => { + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!("AOF sync failed (seq {}, always): {}", manifest.seq, e); + write_error = true; + } + } + FsyncPolicy::EverySec => { + if last_fsync.elapsed() >= std::time::Duration::from_secs(1) { + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!( + "AOF sync failed (seq {}, everysec): {}", + manifest.seq, e + ); + // Non-fatal for everysec: retry next interval + } else { + last_fsync = Instant::now(); + } + } + } + FsyncPolicy::No => {} } } Ok(AofMessage::Shutdown) | Err(_) => { - let _ = file.flush(); - info!("AOF writer shutting down (monoio)"); + if !write_error { + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!("AOF final sync failed (seq {}): {}", manifest.seq, e); + } + } + info!("AOF writer shutting down (monoio, seq {})", manifest.seq); break; } - Ok(AofMessage::Rewrite(_db)) => { - // AOF rewrite under monoio: not yet implemented + Ok(AofMessage::Rewrite(db)) => { + if !write_error { + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!("AOF pre-rewrite sync failed (seq {}): {}", manifest.seq, e); + } + } + match do_rewrite_single(&db, &mut manifest, &mut file) { + Ok(()) => { + write_error = false; // Reset on successful rewrite + } + Err(e) => error!("AOF rewrite failed (seq {}): {}", manifest.seq, e), + } + } + Ok(AofMessage::RewriteSharded(shard_dbs)) => { + if !write_error { + if let Err(e) = file.flush().and_then(|_| file.sync_data()) { + error!("AOF pre-rewrite sync failed (seq {}): {}", manifest.seq, e); + } + } + match do_rewrite_sharded(&shard_dbs, &mut manifest, &mut file) { + Ok(()) => { + write_error = false; + } + Err(e) => error!("AOF rewrite failed (seq {}): {}", manifest.seq, e), + } } } } @@ -190,6 +283,19 @@ pub async fn aof_writer_task( } } } + Ok(AofMessage::RewriteSharded(shard_dbs)) => { + let _ = writer.flush().await; + let _ = writer.get_ref().sync_data().await; + if let Err(e) = rewrite_aof_sharded_sync(&shard_dbs, &aof_path) { + error!("AOF rewrite (sharded) failed: {}", e); + } + let reopen_result: Result = tokio::fs::OpenOptions::new() + .create(true).append(true).open(&aof_path).await; + match reopen_result { + Ok(f) => writer = tokio::io::BufWriter::new(f), + Err(e) => { error!("Failed to reopen AOF after rewrite: {}", e); return; } + } + } Ok(AofMessage::Shutdown) | Err(_) => { let _ = writer.flush().await; let _ = writer.get_ref().sync_data().await; @@ -233,8 +339,35 @@ pub fn replay_aof( return Ok(0); } - let total_len = data.len(); - let mut buf = BytesMut::from(&data[..]); + // Detect RDB preamble: if the file starts with "MOON" magic, load the binary + // RDB section first, then replay any RESP commands appended after it. + let (rdb_keys, resp_start) = if data.starts_with(b"MOON") { + match crate::persistence::rdb::load_from_bytes(databases, &data) { + Ok((keys, consumed)) => { + info!( + "AOF RDB preamble loaded: {} keys ({} bytes)", + keys, consumed + ); + (keys, consumed) + } + Err(e) => { + // Data starts with MOON magic — it IS RDB format. + // Falling back to RESP would parse garbage. Propagate the error. + return Err(e); + } + } + } else { + (0, 0) + }; + + // If the entire file was RDB (no RESP tail), we're done + if resp_start >= data.len() { + return Ok(rdb_keys); + } + + let resp_data = &data[resp_start..]; + let total_len = resp_data.len(); + let mut buf = BytesMut::from(resp_data); let config = ParseConfig::default(); let mut selected_db: usize = 0; let mut count: usize = 0; @@ -314,12 +447,13 @@ pub fn replay_aof( ); } - Ok(count) + Ok(rdb_keys + count) } /// Generate synthetic RESP commands from the current database state for AOF rewriting. /// /// Produces commands for all 5 data types plus PEXPIRE for keys with TTL. +#[allow(dead_code)] // Retained for RESP-only AOF rewrite fallback and testing pub fn generate_rewrite_commands(databases: &[Database]) -> BytesMut { let mut buf = BytesMut::new(); let now_ms = current_time_ms(); @@ -518,12 +652,11 @@ pub fn generate_rewrite_commands(databases: &[Database]) -> BytesMut { buf } -/// Rewrite the AOF file with synthetic commands from current database state. +/// Snapshot databases and generate compacted AOF commands. /// -/// Writes to a temporary file first, then atomically renames for crash safety. -#[cfg(feature = "runtime-tokio")] -pub async fn rewrite_aof(db: SharedDatabases, aof_path: &Path) -> Result<(), MoonError> { - // Clone database state: lock each db individually with read lock +/// Shared by both the async (tokio) and sync (monoio) rewrite paths. +#[allow(dead_code)] +fn snapshot_and_generate(db: &SharedDatabases) -> BytesMut { let snapshot: Vec<(Vec<(CompactKey, Entry)>, u32)> = db .iter() .map(|lock| { @@ -538,7 +671,6 @@ pub async fn rewrite_aof(db: SharedDatabases, aof_path: &Path) -> Result<(), Moo }) .collect(); - // Reconstruct temporary Database objects for generate_rewrite_commands let mut temp_dbs: Vec = Vec::with_capacity(snapshot.len()); for (entries, _base_ts) in &snapshot { let mut db = Database::new(); @@ -548,11 +680,130 @@ pub async fn rewrite_aof(db: SharedDatabases, aof_path: &Path) -> Result<(), Moo temp_dbs.push(db); } - let commands = generate_rewrite_commands(&temp_dbs); + generate_rewrite_commands(&temp_dbs) +} + +/// Multi-part rewrite: snapshot single-shard databases → RDB base → advance manifest. +#[cfg(feature = "runtime-monoio")] +fn do_rewrite_single( + db: &SharedDatabases, + manifest: &mut crate::persistence::aof_manifest::AofManifest, + file: &mut std::fs::File, +) -> Result<(), MoonError> { + // Capture (entries, base_ts) per database — preserves original base_ts for correct TTL. + let snapshot: Vec<( + Vec<( + crate::storage::compact_key::CompactKey, + crate::storage::entry::Entry, + )>, + u32, + )> = db + .iter() + .map(|lock| { + let guard = lock.read(); + let base_ts = guard.base_timestamp(); + let entries = guard + .data() + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + (entries, base_ts) + }) + .collect(); + + let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&snapshot)?; + let new_incr = manifest.advance(&rdb_bytes)?; + + // Switch writer to new incr file + *file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&new_incr) + .map_err(|e| AofError::Io { + path: new_incr, + source: e, + })?; + + Ok(()) +} + +/// Multi-part rewrite: snapshot all shards → merged RDB base → advance manifest. +#[cfg(feature = "runtime-monoio")] +fn do_rewrite_sharded( + shard_dbs: &crate::shard::shared_databases::ShardDatabases, + manifest: &mut crate::persistence::aof_manifest::AofManifest, + file: &mut std::fs::File, +) -> Result<(), MoonError> { + // Capture (entries, base_ts) per merged database, preserving original base_ts for TTL. + let db_count = shard_dbs.db_count(); + let mut merged: Vec<( + Vec<( + crate::storage::compact_key::CompactKey, + crate::storage::entry::Entry, + )>, + u32, + )> = (0..db_count).map(|_| (Vec::new(), 0u32)).collect(); + + for shard_locks in shard_dbs.all_shard_dbs() { + for (db_idx, lock) in shard_locks.iter().enumerate() { + let guard = lock.read(); + let base_ts = guard.base_timestamp(); + let now_ms = current_time_ms(); + if merged[db_idx].0.is_empty() { + merged[db_idx].1 = base_ts; + } + for (key, entry) in guard.data().iter() { + if !entry.is_expired_at(base_ts, now_ms) { + merged[db_idx].0.push((key.clone(), entry.clone())); + } + } + } + } + + let rdb_bytes = crate::persistence::rdb::save_snapshot_to_bytes(&merged)?; + let new_incr = manifest.advance(&rdb_bytes)?; + + *file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&new_incr) + .map_err(|e| AofError::Io { + path: new_incr, + source: e, + })?; + + Ok(()) +} + +/// Rewrite the AOF file with RDB preamble (binary base + empty RESP incremental). +/// +/// Uses the same strategy as Redis 7+ `aof-use-rdb-preamble yes`: +/// the rewritten AOF starts with a full RDB snapshot (compact binary), +/// and new writes are appended as RESP after it. On startup, the loader +/// detects the RDB magic and reads the binary preamble, then switches +/// to RESP parsing for any incremental commands appended after. +#[allow(dead_code)] // Retained for legacy single-file and tokio path +fn rewrite_aof_sync(db: &SharedDatabases, aof_path: &Path) -> Result<(), MoonError> { + // Snapshot under read locks, build temp Database objects for RDB serialization + let snapshot: Vec = db + .iter() + .map(|lock| { + let guard = lock.read(); + let mut temp = Database::new(); + let now_ms = current_time_ms(); + for (k, v) in guard.data().iter() { + if !v.is_expired_at(guard.base_timestamp(), now_ms) { + temp.set(k.to_bytes(), v.clone()); + } + } + temp + }) + .collect(); + + let rdb_bytes = crate::persistence::rdb::save_to_bytes(&snapshot)?; - // Write to temp file, then atomic rename let tmp_path = aof_path.with_extension("aof.tmp"); - std::fs::write(&tmp_path, &commands).map_err(|e| AofError::Io { + std::fs::write(&tmp_path, &rdb_bytes).map_err(|e| AofError::Io { path: tmp_path.clone(), source: e, })?; @@ -565,10 +816,77 @@ pub async fn rewrite_aof(db: SharedDatabases, aof_path: &Path) -> Result<(), Moo ), })?; - info!("AOF rewrite complete: {} bytes", commands.len()); + info!( + "AOF rewrite complete (RDB preamble): {} bytes", + rdb_bytes.len() + ); Ok(()) } +/// Rewrite the AOF in sharded mode with RDB preamble. +/// +/// Merges all shards' databases into a single RDB snapshot, writes it as +/// the AOF base file. New incremental writes are appended as RESP after. +#[allow(dead_code)] +fn rewrite_aof_sharded_sync( + shard_dbs: &crate::shard::shared_databases::ShardDatabases, + aof_path: &Path, +) -> Result<(), MoonError> { + let db_count = shard_dbs.db_count(); + let now_ms = current_time_ms(); + let mut merged_dbs: Vec = (0..db_count).map(|_| Database::new()).collect(); + + for shard_locks in shard_dbs.all_shard_dbs() { + for (db_idx, lock) in shard_locks.iter().enumerate() { + let guard = lock.read(); + for (key, entry) in guard.data().iter() { + if !entry.is_expired_at(guard.base_timestamp(), now_ms) { + merged_dbs[db_idx].set(key.to_bytes(), entry.clone()); + } + } + } + } + + let rdb_bytes = crate::persistence::rdb::save_to_bytes(&merged_dbs)?; + + let tmp_path = aof_path.with_extension("aof.tmp"); + std::fs::write(&tmp_path, &rdb_bytes).map_err(|e| AofError::Io { + path: tmp_path.clone(), + source: e, + })?; + std::fs::rename(&tmp_path, aof_path).map_err(|e| AofError::RewriteFailed { + detail: format!( + "rename {} -> {}: {}", + tmp_path.display(), + aof_path.display(), + e + ), + })?; + + info!( + "AOF rewrite (sharded, RDB preamble) complete: {} bytes", + rdb_bytes.len() + ); + Ok(()) +} + +/// Reopen AOF file in append mode after atomic rewrite replaced it. +#[allow(dead_code)] +fn reopen_aof_sync(aof_path: &Path) -> Result { + std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(aof_path) +} + +/// Rewrite the AOF file (tokio async wrapper). +/// +/// Delegates to `rewrite_aof_sync` — the actual I/O is synchronous (temp write + rename). +#[cfg(feature = "runtime-tokio")] +pub async fn rewrite_aof(db: SharedDatabases, aof_path: &Path) -> Result<(), MoonError> { + rewrite_aof_sync(&db, aof_path) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/persistence/aof_manifest.rs b/src/persistence/aof_manifest.rs new file mode 100644 index 000000000..84b884a8f --- /dev/null +++ b/src/persistence/aof_manifest.rs @@ -0,0 +1,314 @@ +//! Multi-part AOF manifest: tracks base (RDB) and incremental (RESP) files. +//! +//! Implements the same directory-based AOF format as Redis 7+: +//! ```text +//! appendonlydir/ +//! moon.aof.1.base.rdb # RDB snapshot base +//! moon.aof.1.incr.aof # Incremental RESP since base +//! moon.aof.manifest # This file +//! ``` +//! +//! The manifest is a simple text file listing the active base and incremental +//! files with their sequence numbers. On BGREWRITEAOF, the sequence increments, +//! a new base + incr pair is created, and old files are deleted. + +use std::io::Write; +use std::path::{Path, PathBuf}; + +use tracing::{error, info, warn}; + +const MANIFEST_NAME: &str = "moon.aof.manifest"; +const AOF_DIR_NAME: &str = "appendonlydir"; + +/// Active AOF file set tracked by the manifest. +#[derive(Debug, Clone)] +pub struct AofManifest { + /// Base directory (parent of `appendonlydir/`) + pub dir: PathBuf, + /// Current sequence number (incremented on each rewrite) + pub seq: u64, +} + +impl AofManifest { + /// Path to the `appendonlydir/` directory. + pub fn aof_dir(&self) -> PathBuf { + self.dir.join(AOF_DIR_NAME) + } + + /// Path to the manifest file. + pub fn manifest_path(&self) -> PathBuf { + self.aof_dir().join(MANIFEST_NAME) + } + + /// Path to the base RDB file for the current sequence. + pub fn base_path(&self) -> PathBuf { + self.aof_dir() + .join(format!("moon.aof.{}.base.rdb", self.seq)) + } + + /// Path to the incremental RESP file for the current sequence. + pub fn incr_path(&self) -> PathBuf { + self.aof_dir() + .join(format!("moon.aof.{}.incr.aof", self.seq)) + } + + /// Path to the base RDB file for a given sequence. + pub fn base_path_seq(&self, seq: u64) -> PathBuf { + self.aof_dir().join(format!("moon.aof.{}.base.rdb", seq)) + } + + /// Path to the incremental RESP file for a given sequence. + pub fn incr_path_seq(&self, seq: u64) -> PathBuf { + self.aof_dir().join(format!("moon.aof.{}.incr.aof", seq)) + } + + /// Create the `appendonlydir/` and write the initial manifest. + pub fn initialize(dir: &Path) -> std::io::Result { + let manifest = Self { + dir: dir.to_path_buf(), + seq: 1, + }; + std::fs::create_dir_all(manifest.aof_dir())?; + manifest.write_manifest()?; + Ok(manifest) + } + + /// Load manifest from disk. Returns `None` if manifest doesn't exist. + pub fn load(dir: &Path) -> Option { + let aof_dir = dir.join(AOF_DIR_NAME); + let manifest_path = aof_dir.join(MANIFEST_NAME); + + if !manifest_path.exists() { + return None; + } + + let content = match std::fs::read_to_string(&manifest_path) { + Ok(c) => c, + Err(e) => { + error!("Failed to read AOF manifest: {}", e); + return None; + } + }; + + let mut seq = 0u64; + for line in content.lines() { + let line = line.trim(); + if let Some(val) = line.strip_prefix("seq ") { + if let Ok(n) = val.parse::() { + seq = n; + } + } + } + + if seq == 0 { + error!("AOF manifest has no valid sequence number"); + return None; + } + + Some(Self { + dir: dir.to_path_buf(), + seq, + }) + } + + /// Write the manifest file atomically (write tmp + rename). + pub fn write_manifest(&self) -> std::io::Result<()> { + let manifest_path = self.manifest_path(); + let tmp_path = manifest_path.with_extension("tmp"); + + let content = format!( + "seq {}\nbase moon.aof.{}.base.rdb\nincr moon.aof.{}.incr.aof\n", + self.seq, self.seq, self.seq + ); + + let mut f = std::fs::File::create(&tmp_path)?; + f.write_all(content.as_bytes())?; + f.sync_data()?; + std::fs::rename(&tmp_path, &manifest_path)?; + Ok(()) + } + + /// Advance to the next sequence: write new base RDB, create new incr file, + /// update manifest, delete old files. + /// + /// Returns the path to the new incremental file (caller should switch writing to it). + pub fn advance(&mut self, rdb_bytes: &[u8]) -> Result { + let old_seq = self.seq; + let new_seq = old_seq + 1; + + let aof_dir = self.aof_dir(); + std::fs::create_dir_all(&aof_dir).map_err(|e| crate::error::AofError::Io { + path: aof_dir.clone(), + source: e, + })?; + + // 1. Write new base RDB (atomic: tmp + rename) + let new_base = self.base_path_seq(new_seq); + let tmp_base = new_base.with_extension("rdb.tmp"); + std::fs::write(&tmp_base, rdb_bytes).map_err(|e| crate::error::AofError::Io { + path: tmp_base.clone(), + source: e, + })?; + std::fs::rename(&tmp_base, &new_base).map_err(|e| { + crate::error::AofError::RewriteFailed { + detail: format!("rename base: {}", e), + } + })?; + + // 2. Create empty new incremental file + let new_incr = self.incr_path_seq(new_seq); + std::fs::File::create(&new_incr).map_err(|e| crate::error::AofError::Io { + path: new_incr.clone(), + source: e, + })?; + + // 3. Update manifest (atomic) + self.seq = new_seq; + self.write_manifest() + .map_err(|e| crate::error::AofError::Io { + path: self.manifest_path(), + source: e, + })?; + + // 4. Delete old files (best-effort) + let old_base = self.base_path_seq(old_seq); + let old_incr = self.incr_path_seq(old_seq); + if old_base.exists() { + if let Err(e) = std::fs::remove_file(&old_base) { + warn!("Failed to delete old base {}: {}", old_base.display(), e); + } + } + if old_incr.exists() { + if let Err(e) = std::fs::remove_file(&old_incr) { + warn!("Failed to delete old incr {}: {}", old_incr.display(), e); + } + } + + info!( + "AOF advanced to seq {}: base={} bytes, incr={}", + new_seq, + rdb_bytes.len(), + new_incr.display() + ); + + Ok(new_incr) + } +} + +/// Replay multi-part AOF: load base RDB then replay incremental RESP. +/// +/// Returns total keys/commands loaded. +pub fn replay_multi_part( + databases: &mut [crate::storage::Database], + manifest: &AofManifest, + engine: &dyn crate::persistence::replay::CommandReplayEngine, +) -> Result { + let mut total = 0usize; + + // Load base RDB + let base_path = manifest.base_path(); + if base_path.exists() { + match crate::persistence::rdb::load(databases, &base_path) { + Ok(n) => { + info!( + "AOF base RDB loaded: {} keys from {}", + n, + base_path.display() + ); + total += n; + } + Err(e) => { + // Base RDB is corrupt or unreadable — applying incremental + // deltas on top of missing/corrupt base gives wrong results. + error!("AOF base RDB load failed: {}", e); + return Err(e); + } + } + } else { + warn!("AOF base RDB not found: {}", base_path.display()); + } + + // Replay incremental RESP + let incr_path = manifest.incr_path(); + if incr_path.exists() { + let data = std::fs::read(&incr_path)?; + if !data.is_empty() { + // Pure RESP — use replay_aof_resp (no RDB preamble detection needed) + let count = replay_incr_resp(databases, &data, engine)?; + info!( + "AOF incr replayed: {} commands from {}", + count, + incr_path.display() + ); + total += count; + } + } + + Ok(total) +} + +/// Replay pure RESP commands from a byte slice. +fn replay_incr_resp( + databases: &mut [crate::storage::Database], + data: &[u8], + engine: &dyn crate::persistence::replay::CommandReplayEngine, +) -> Result { + use crate::protocol::{Frame, ParseConfig, parse}; + use bytes::BytesMut; + + let total_len = data.len(); + let mut buf = BytesMut::from(data); + let config = ParseConfig::default(); + let mut selected_db: usize = 0; + let mut count: usize = 0; + + loop { + if buf.is_empty() { + break; + } + match parse::parse(&mut buf, &config) { + Ok(Some(frame)) => { + let (cmd, cmd_args) = match &frame { + Frame::Array(arr) if !arr.is_empty() => { + let name = match &arr[0] { + Frame::BulkString(s) => s.as_ref(), + Frame::SimpleString(s) => s.as_ref(), + _ => { + count += 1; + continue; + } + }; + (name as &[u8], &arr[1..]) + } + _ => { + count += 1; + continue; + } + }; + engine.replay_command(databases, cmd, cmd_args, &mut selected_db); + count += 1; + } + Ok(None) => { + if !buf.is_empty() { + let offset = total_len - buf.len(); + warn!( + "AOF incr truncated: {} bytes at offset {}", + buf.len(), + offset + ); + } + break; + } + Err(_) => { + let _ = buf.split_to(1); + if let Some(pos) = buf.iter().position(|&b| b == b'*') { + let _ = buf.split_to(pos); + } else { + break; + } + } + } + } + + Ok(count) +} diff --git a/src/persistence/mod.rs b/src/persistence/mod.rs index 907d104f3..c63438de3 100644 --- a/src/persistence/mod.rs +++ b/src/persistence/mod.rs @@ -1,4 +1,5 @@ pub mod aof; +pub mod aof_manifest; pub mod auto_save; pub mod rdb; pub mod redis_rdb; diff --git a/src/persistence/rdb.rs b/src/persistence/rdb.rs index a14ec0d6b..dbc8a3bad 100644 --- a/src/persistence/rdb.rs +++ b/src/persistence/rdb.rs @@ -47,7 +47,11 @@ const EOF_MARKER: u8 = 0xFF; /// Uses atomic write (write to .tmp, then rename) for crash safety. /// Expired keys are skipped. Empty databases are skipped. /// Footer contains CRC32 checksum of all preceding bytes. -pub fn save(databases: &[Database], path: &Path) -> Result<(), MoonError> { +/// Serialize all databases to RDB format in memory. +/// +/// Returns the complete RDB byte stream (header + entries + footer + CRC32). +/// Used by both `save()` (file) and AOF RDB-preamble rewrite. +pub fn save_to_bytes(databases: &[Database]) -> Result, MoonError> { let mut buf = Vec::new(); // Header @@ -60,7 +64,6 @@ pub fn save(databases: &[Database], path: &Path) -> Result<(), MoonError> { for (db_idx, db) in databases.iter().enumerate() { let base_ts = db.base_timestamp(); let data = db.data(); - // Collect non-expired entries let live: Vec<_> = data .iter() .filter(|(_, entry)| !entry.is_expired_at(base_ts, now_ms)) @@ -86,6 +89,12 @@ pub fn save(databases: &[Database], path: &Path) -> Result<(), MoonError> { let checksum = hasher.finalize(); buf.write_all(&checksum.to_le_bytes())?; + Ok(buf) +} + +pub fn save(databases: &[Database], path: &Path) -> Result<(), MoonError> { + let buf = save_to_bytes(databases)?; + // Atomic write: write to tmp, then rename let tmp_path = path.with_extension("rdb.tmp"); std::fs::write(&tmp_path, &buf).map_err(|e| RdbError::Io { @@ -154,6 +163,46 @@ pub fn save_from_snapshot( Ok(()) } +/// Serialize snapshot data (with correct base_ts per database) to RDB bytes in memory. +/// +/// Unlike `save_to_bytes(&[Database])` which reads base_ts from each Database, +/// this takes explicit (entries, base_ts) tuples — critical for AOF rewrite where +/// entries are cloned into temporary storage and the original base_ts must be preserved. +pub fn save_snapshot_to_bytes( + snapshot: &[(Vec<(CompactKey, Entry)>, u32)], +) -> Result, MoonError> { + let mut buf = Vec::new(); + + buf.write_all(RDB_MAGIC)?; + buf.write_all(&[RDB_VERSION])?; + + let now_ms = current_time_ms(); + + for (db_idx, (entries, base_ts)) in snapshot.iter().enumerate() { + let live: Vec<_> = entries + .iter() + .filter(|(_, e)| !e.is_expired_at(*base_ts, now_ms)) + .collect(); + if live.is_empty() { + continue; + } + + buf.write_all(&[DB_SELECTOR])?; + buf.write_all(&[db_idx as u8])?; + + for (key, entry) in live { + write_entry(&mut buf, key.as_bytes(), entry, *base_ts)?; + } + } + + buf.write_all(&[EOF_MARKER])?; + let mut hasher = Hasher::new(); + hasher.update(&buf); + buf.write_all(&hasher.finalize().to_le_bytes())?; + + Ok(buf) +} + /// Load an RDB file and populate databases. Returns total keys loaded. /// /// On any error (missing file, corrupt data, bad checksum), returns Err. @@ -175,25 +224,28 @@ pub fn load(databases: &mut [Database], path: &Path) -> Result .into()); } + // Wrap in Bytes for zero-copy slicing (shared refcount, no copy) + let shared_buf = Bytes::from(data); + // Verify CRC32: all bytes except last 4 vs last 4 bytes - let (payload, checksum_bytes) = data.split_at(data.len() - 4); + let payload_len = shared_buf.len() - 4; let stored_checksum = u32::from_le_bytes([ - checksum_bytes[0], - checksum_bytes[1], - checksum_bytes[2], - checksum_bytes[3], + shared_buf[payload_len], + shared_buf[payload_len + 1], + shared_buf[payload_len + 2], + shared_buf[payload_len + 3], ]); let mut hasher = Hasher::new(); - hasher.update(payload); + hasher.update(&shared_buf[..payload_len]); let computed_checksum = hasher.finalize(); if stored_checksum != computed_checksum { return Err(RdbError::ChecksumMismatch.into()); } - let mut cursor = Cursor::new(payload); + let mut cursor = Cursor::new(&shared_buf[..payload_len] as &[u8]); // Verify magic - let mut magic = [0u8; 4]; // "MOON" is 4 bytes + let mut magic = [0u8; 4]; cursor.read_exact(&mut magic).map_err(|e| RdbError::Io { path: path.to_path_buf(), source: e, @@ -218,18 +270,28 @@ pub fn load(databases: &mut [Database], path: &Path) -> Result .into()); } + // Cache timestamps once (Fix #4: avoid syscall per entry) let now_ms = current_time_ms(); + let now_secs = (now_ms / 1000) as u32; + + // First pass: count entries per database for pre-sizing (Fix #2) + let entry_counts = count_entries_per_db(&cursor, databases.len()); + + // Pre-size DashTables to avoid segment splits during load (Fix #2) + for (db_idx, &count) in entry_counts.iter().enumerate() { + if count > 0 && db_idx < databases.len() { + databases[db_idx].reserve(count); + } + } let mut total_keys = 0usize; - let mut skipped_entries = 0usize; let mut current_db: usize = 0; loop { let mut tag = [0u8; 1]; if cursor.read_exact(&mut tag).is_err() { - // Truncated tail: no more data to read, treat as implicit EOF tracing::warn!( - "RDB load: truncated tail after {} keys (no EOF marker), treating as end of file", + "RDB load: truncated tail after {} keys (no EOF marker)", total_keys ); break; @@ -256,33 +318,22 @@ pub fn load(databases: &mut [Database], path: &Path) -> Result } } type_tag => { - // Mid-stream corruption recovery: log+skip on entry parse failure - match read_entry(&mut cursor, type_tag) { + match read_entry_zero_copy(&mut cursor, type_tag, &shared_buf, now_secs) { Ok((key, entry)) => { - // Skip entries whose TTL is already in the past - if entry.has_expiry() && entry.is_expired_at(current_secs(), now_ms) { + if entry.has_expiry() && entry.is_expired_at(now_secs, now_ms) { continue; } if current_db < databases.len() { - databases[current_db].set(key, entry); + // Fix #3: skip duplicate check + memory accounting + databases[current_db].insert_for_load(key, entry); total_keys += 1; } } Err(e) => { - let offset = cursor.position(); - tracing::warn!( - "RDB load: skipping corrupted entry at offset {}: {}", - offset, - e - ); - skipped_entries += 1; - // Cannot reliably skip to next entry in a variable-length - // format without framing, so break out of the loop. - // Entries loaded so far are valid (checksum passed). tracing::warn!( - "RDB load: stopping mid-stream recovery after {} skipped entries; \ - {} keys loaded successfully", - skipped_entries, + "RDB load: corrupted entry at offset {}: {}. {} keys loaded.", + cursor.position(), + e, total_keys ); break; @@ -292,17 +343,525 @@ pub fn load(databases: &mut [Database], path: &Path) -> Result } } - if skipped_entries > 0 { - tracing::warn!( - "RDB load completed with {} entries skipped due to corruption, {} keys loaded", - skipped_entries, - total_keys - ); + // Fix #3: single-pass memory recalculation after all inserts + for db in databases.iter_mut() { + db.recalculate_memory(); } Ok(total_keys) } +/// Fast first-pass: count entries per database without parsing values. +/// Scans type tags and skips over entry payloads to count keys per db_idx. +fn count_entries_per_db(cursor: &Cursor<&[u8]>, db_count: usize) -> Vec { + let mut counts = vec![0usize; db_count]; + let data = cursor.get_ref(); + let mut pos = cursor.position() as usize; + let mut current_db = 0usize; + + while pos < data.len() { + let tag = data[pos]; + pos += 1; + + match tag { + EOF_MARKER => break, + DB_SELECTOR => { + if pos < data.len() { + current_db = data[pos] as usize; + pos += 1; + } else { + break; + } + } + TYPE_STRING | TYPE_HASH | TYPE_LIST | TYPE_SET | TYPE_SORTED_SET | TYPE_STREAM => { + if current_db < db_count { + counts[current_db] += 1; + } + // Skip over the entry payload without parsing + if let Some(new_pos) = skip_entry(data, pos, tag) { + pos = new_pos; + } else { + break; + } + } + _ => break, + } + } + + counts +} + +/// Skip over an RDB entry's bytes without allocating or parsing values. +/// Returns the new position after the entry, or None if data is truncated. +fn skip_entry(data: &[u8], mut pos: usize, type_tag: u8) -> Option { + // Skip key + pos = skip_bytes_field(data, pos)?; + // Skip TTL (8 bytes) + pos = pos.checked_add(8)?; + if pos > data.len() { + return None; + } + + match type_tag { + TYPE_STRING => { + pos = skip_bytes_field(data, pos)?; + } + TYPE_HASH => { + let count = read_u32_raw(data, pos)?; + pos += 4; + for _ in 0..count { + pos = skip_bytes_field(data, pos)?; // field + pos = skip_bytes_field(data, pos)?; // value + } + } + TYPE_LIST | TYPE_SET => { + let count = read_u32_raw(data, pos)?; + pos += 4; + for _ in 0..count { + pos = skip_bytes_field(data, pos)?; + } + } + TYPE_SORTED_SET => { + let count = read_u32_raw(data, pos)?; + pos += 4; + for _ in 0..count { + pos = skip_bytes_field(data, pos)?; // member + pos = pos.checked_add(8)?; // f64 score + if pos > data.len() { + return None; + } + } + } + TYPE_STREAM => { + // entry_count(8) + last_id(16) + pos = pos.checked_add(24)?; + if pos > data.len() { + return None; + } + let entry_count = + u64::from_le_bytes(data[pos - 24..pos - 16].try_into().ok()?) as usize; + for _ in 0..entry_count { + pos = pos.checked_add(16)?; // StreamId (ms + seq) + if pos > data.len() { + return None; + } + let field_count = read_u32_raw(data, pos)?; + pos += 4; + for _ in 0..field_count { + pos = skip_bytes_field(data, pos)?; + pos = skip_bytes_field(data, pos)?; + } + } + // Consumer groups + let group_count = read_u32_raw(data, pos)?; + pos += 4; + for _ in 0..group_count { + pos = skip_bytes_field(data, pos)?; // group name + pos = pos.checked_add(16)?; // last_delivered_id + if pos > data.len() { + return None; + } + let pel_count = read_u32_raw(data, pos)?; + pos += 4; + for _ in 0..pel_count { + pos = pos.checked_add(16)?; // StreamId + if pos > data.len() { + return None; + } + pos = skip_bytes_field(data, pos)?; // consumer name + pos = pos.checked_add(16)?; // delivery_time + delivery_count + if pos > data.len() { + return None; + } + } + let consumer_count = read_u32_raw(data, pos)?; + pos += 4; + for _ in 0..consumer_count { + pos = skip_bytes_field(data, pos)?; // consumer name + pos = pos.checked_add(8)?; // seen_time + if pos > data.len() { + return None; + } + let pending_count = read_u32_raw(data, pos)?; + pos += 4; + for _ in 0..pending_count { + pos = pos.checked_add(16)?; // StreamId + if pos > data.len() { + return None; + } + } + } + } + } + _ => return None, + } + + Some(pos) +} + +/// Read u32 LE from raw bytes without cursor overhead. +#[inline] +fn read_u32_raw(data: &[u8], pos: usize) -> Option { + if pos + 4 > data.len() { + return None; + } + Some(u32::from_le_bytes(data[pos..pos + 4].try_into().ok()?) as usize) +} + +/// Skip a length-prefixed bytes field (4-byte LE length + payload). +#[inline] +fn skip_bytes_field(data: &[u8], pos: usize) -> Option { + let len = read_u32_raw(data, pos)?; + let new_pos = pos.checked_add(4)?.checked_add(len)?; + if new_pos > data.len() { + None + } else { + Some(new_pos) + } +} + +/// Zero-copy variant of read_entry: uses shared Bytes buffer and cached timestamps. +fn read_entry_zero_copy( + cursor: &mut Cursor<&[u8]>, + type_tag: u8, + _shared_buf: &Bytes, + cached_secs: u32, +) -> Result<(Bytes, Entry), MoonError> { + let key = read_bytes(cursor)?; + + let mut ttl_buf = [0u8; 8]; + cursor.read_exact(&mut ttl_buf)?; + let ttl_ms = i64::from_le_bytes(ttl_buf); + let expires_at_ms = if ttl_ms > 0 { ttl_ms as u64 } else { 0 }; + + let value = match type_tag { + TYPE_STRING => { + // Fast path: build CompactValue directly from Vec, skipping RedisValue intermediate. + // This avoids: Vec → Bytes → RedisValue::String → from_redis_value → heap_string_vec + // and instead does: Vec → CompactValue directly (one Box alloc, zero copy). + let vec = read_bytes_vec(cursor)?; + let cv = if vec.len() <= 12 { + crate::storage::compact_value::CompactValue::from_redis_value(RedisValue::String( + Bytes::from(vec), + )) + } else { + crate::storage::compact_value::CompactValue::heap_string_vec_direct(vec) + }; + let mut entry = Entry::new_string(Bytes::new()); + entry.value = cv; + if expires_at_ms > 0 { + entry.set_expires_at_ms(cached_secs, expires_at_ms); + } + entry.set_last_access(cached_secs); + entry.set_access_counter(5); + return Ok((key, entry)); + } + TYPE_HASH => { + let count = read_u32(cursor)? as usize; + validate_count(cursor, count, 8, "hash")?; + let mut map = HashMap::with_capacity(count); + for _ in 0..count { + let field = read_bytes(cursor)?; + let val = read_bytes(cursor)?; + map.insert(field, val); + } + RedisValue::Hash(map) + } + TYPE_LIST => { + let count = read_u32(cursor)? as usize; + validate_count(cursor, count, 4, "list")?; + let mut list = VecDeque::with_capacity(count); + for _ in 0..count { + list.push_back(read_bytes(cursor)?); + } + RedisValue::List(list) + } + TYPE_SET => { + let count = read_u32(cursor)? as usize; + validate_count(cursor, count, 4, "set")?; + let mut set = HashSet::with_capacity(count); + for _ in 0..count { + set.insert(read_bytes(cursor)?); + } + RedisValue::Set(set) + } + TYPE_SORTED_SET => { + let count = read_u32(cursor)? as usize; + validate_count(cursor, count, 12, "sorted_set")?; + let mut members = HashMap::with_capacity(count); + let mut tree = BPTree::new(); + for _ in 0..count { + let member = read_bytes(cursor)?; + let mut score_buf = [0u8; 8]; + cursor.read_exact(&mut score_buf)?; + let score = f64::from_le_bytes(score_buf); + members.insert(member.clone(), score); + tree.insert(OrderedFloat(score), member); + } + RedisValue::SortedSetBPTree { tree, members } + } + TYPE_STREAM => { + // Stream parsing: reuse read_bytes (not zero-copy for this rare type) + let mut entry_count_buf = [0u8; 8]; + cursor.read_exact(&mut entry_count_buf)?; + let entry_count = u64::from_le_bytes(entry_count_buf) as usize; + let mut last_id_ms_buf = [0u8; 8]; + let mut last_id_seq_buf = [0u8; 8]; + cursor.read_exact(&mut last_id_ms_buf)?; + cursor.read_exact(&mut last_id_seq_buf)?; + let last_id = StreamId { + ms: u64::from_le_bytes(last_id_ms_buf), + seq: u64::from_le_bytes(last_id_seq_buf), + }; + let mut stream = StreamData::new(); + stream.last_id = last_id; + validate_count(cursor, entry_count, 20, "stream_entries")?; + for _ in 0..entry_count { + let mut ms_buf = [0u8; 8]; + let mut seq_buf = [0u8; 8]; + cursor.read_exact(&mut ms_buf)?; + cursor.read_exact(&mut seq_buf)?; + let id = StreamId { + ms: u64::from_le_bytes(ms_buf), + seq: u64::from_le_bytes(seq_buf), + }; + let field_count = read_u32(cursor)? as usize; + validate_count(cursor, field_count, 8, "stream_fields")?; + let mut fields = Vec::with_capacity(field_count); + for _ in 0..field_count { + fields.push((read_bytes(cursor)?, read_bytes(cursor)?)); + } + stream.entries.insert(id, fields); + stream.length += 1; + } + let group_count = read_u32(cursor)? as usize; + for _ in 0..group_count { + let group_name = read_bytes(cursor)?; + let mut gld_ms = [0u8; 8]; + let mut gld_seq = [0u8; 8]; + cursor.read_exact(&mut gld_ms)?; + cursor.read_exact(&mut gld_seq)?; + let last_delivered_id = StreamId { + ms: u64::from_le_bytes(gld_ms), + seq: u64::from_le_bytes(gld_seq), + }; + let pel_count = read_u32(cursor)? as usize; + let mut pel = BTreeMap::new(); + for _ in 0..pel_count { + let mut pid_ms = [0u8; 8]; + let mut pid_seq = [0u8; 8]; + cursor.read_exact(&mut pid_ms)?; + cursor.read_exact(&mut pid_seq)?; + let pid = StreamId { + ms: u64::from_le_bytes(pid_ms), + seq: u64::from_le_bytes(pid_seq), + }; + let consumer_name = read_bytes(cursor)?; + let mut dt_buf = [0u8; 8]; + let mut dc_buf = [0u8; 8]; + cursor.read_exact(&mut dt_buf)?; + cursor.read_exact(&mut dc_buf)?; + pel.insert( + pid, + crate::storage::stream::PendingEntry { + consumer: consumer_name, + delivery_time: u64::from_le_bytes(dt_buf), + delivery_count: u64::from_le_bytes(dc_buf), + }, + ); + } + let consumer_count = read_u32(cursor)? as usize; + let mut consumers = HashMap::new(); + for _ in 0..consumer_count { + let cname = read_bytes(cursor)?; + let mut st_buf = [0u8; 8]; + cursor.read_exact(&mut st_buf)?; + let seen_time = u64::from_le_bytes(st_buf); + let pending_count = read_u32(cursor)? as usize; + let mut pending = BTreeMap::new(); + for _ in 0..pending_count { + let mut cid_ms = [0u8; 8]; + let mut cid_seq = [0u8; 8]; + cursor.read_exact(&mut cid_ms)?; + cursor.read_exact(&mut cid_seq)?; + pending.insert( + StreamId { + ms: u64::from_le_bytes(cid_ms), + seq: u64::from_le_bytes(cid_seq), + }, + (), + ); + } + consumers.insert( + cname.clone(), + crate::storage::stream::Consumer { + name: cname, + pending, + seen_time, + }, + ); + } + stream.groups.insert( + group_name, + crate::storage::stream::ConsumerGroup { + last_delivered_id, + pel, + consumers, + }, + ); + } + RedisValue::Stream(Box::new(stream)) + } + _ => return Err(RdbError::UnsupportedType { type_tag }.into()), + }; + + let mut entry = Entry::new_string(Bytes::new()); + entry.value = crate::storage::compact_value::CompactValue::from_redis_value(value); + if expires_at_ms > 0 { + entry.set_expires_at_ms(cached_secs, expires_at_ms); + } + entry.set_last_access(cached_secs); + entry.set_access_counter(5); + + Ok((key, entry)) +} + +/// Load an RDB snapshot from a byte slice (for AOF RDB-preamble format). +/// +/// Returns `(keys_loaded, bytes_consumed)`. The caller can use `bytes_consumed` +/// to find the start of any RESP commands appended after the RDB preamble. +pub fn load_from_bytes( + databases: &mut [Database], + data: &[u8], +) -> Result<(usize, usize), MoonError> { + if data.len() < RDB_MAGIC.len() + 1 + 1 + 4 { + return Err(RdbError::Corrupted { + detail: "RDB preamble too small".into(), + } + .into()); + } + + // Find EOF_MARKER to determine RDB section length. + // The RDB section is: header + entries + EOF_MARKER(1) + CRC32(4). + // We scan for EOF_MARKER (0xFF) — the first one after the header that's + // immediately followed by a valid CRC32 of the preceding bytes. + let mut rdb_end = None; + // Start scanning after header (MOON + version = 5 bytes) + for i in 5..data.len().saturating_sub(4) { + if data[i] == EOF_MARKER { + let payload = &data[..=i]; // everything up to and including EOF_MARKER + if let Some(checksum_bytes) = data.get(i + 1..i + 5) { + let stored = u32::from_le_bytes([ + checksum_bytes[0], + checksum_bytes[1], + checksum_bytes[2], + checksum_bytes[3], + ]); + let mut hasher = Hasher::new(); + hasher.update(payload); + if hasher.finalize() == stored { + rdb_end = Some(i + 5); // past CRC32 + break; + } + } + } + } + + let rdb_len = rdb_end.ok_or_else(|| { + MoonError::from(RdbError::Corrupted { + detail: "RDB preamble: no valid EOF+CRC found".into(), + }) + })?; + + // Load using the same logic as `load`, but from the byte slice + let payload = &data[..rdb_len - 4]; // exclude CRC32 + let mut cursor = Cursor::new(payload); + + // Skip magic + version + let mut magic = [0u8; 4]; + cursor.read_exact(&mut magic).map_err(|e| RdbError::Io { + path: std::path::PathBuf::from(""), + source: e, + })?; + if &magic != RDB_MAGIC { + return Err(RdbError::Corrupted { + detail: "invalid RDB magic in AOF preamble".into(), + } + .into()); + } + let mut version = [0u8; 1]; + cursor.read_exact(&mut version).map_err(|e| RdbError::Io { + path: std::path::PathBuf::from(""), + source: e, + })?; + if version[0] != RDB_VERSION { + return Err(RdbError::UnsupportedVersion { + version: version[0] as u32, + } + .into()); + } + + let now_ms = current_time_ms(); + let now_secs = (now_ms / 1000) as u32; + let shared_buf = Bytes::copy_from_slice(data); + let mut total_keys = 0usize; + let mut current_db: usize = 0; + + // Pre-size DashTables + let entry_counts = count_entries_per_db(&cursor, databases.len()); + for (db_idx, &count) in entry_counts.iter().enumerate() { + if count > 0 && db_idx < databases.len() { + databases[db_idx].reserve(count); + } + } + + loop { + let mut tag = [0u8; 1]; + if cursor.read_exact(&mut tag).is_err() { + break; + } + match tag[0] { + EOF_MARKER => break, + DB_SELECTOR => { + let mut db_idx = [0u8; 1]; + cursor.read_exact(&mut db_idx).map_err(|e| RdbError::Io { + path: std::path::PathBuf::from(""), + source: e, + })?; + current_db = db_idx[0] as usize; + if current_db >= databases.len() { + return Err(RdbError::Corrupted { + detail: format!( + "RDB preamble references database {} but only {} configured", + current_db, + databases.len() + ), + } + .into()); + } + } + type_tag => match read_entry_zero_copy(&mut cursor, type_tag, &shared_buf, now_secs) { + Ok((key, entry)) => { + if entry.has_expiry() && entry.is_expired_at(now_secs, now_ms) { + continue; + } + if current_db < databases.len() { + databases[current_db].insert_for_load(key, entry); + total_keys += 1; + } + } + Err(_) => break, + }, + } + } + + for db in databases.iter_mut() { + db.recalculate_memory(); + } + + Ok((total_keys, rdb_len)) +} + /// Distribute keys from loaded databases to the correct per-shard databases. /// /// After loading an RDB file into temporary databases, this function routes each key @@ -756,7 +1315,8 @@ fn validate_count( pub(crate) fn read_bytes(cursor: &mut Cursor<&[u8]>) -> Result { let len = read_u32(cursor)? as usize; - let remaining = cursor.get_ref().len() - cursor.position() as usize; + let pos = cursor.position() as usize; + let remaining = cursor.get_ref().len() - pos; if len > remaining { return Err(RdbError::Corrupted { detail: format!( @@ -766,9 +1326,51 @@ pub(crate) fn read_bytes(cursor: &mut Cursor<&[u8]>) -> Result } .into()); } - let mut data = vec![0u8; len]; - cursor.read_exact(&mut data)?; - Ok(Bytes::from(data)) + let slice = &cursor.get_ref()[pos..pos + len]; + cursor.set_position((pos + len) as u64); + Ok(Bytes::copy_from_slice(slice)) +} + +/// Read bytes as owned Vec — avoids Bytes intermediate for RDB load path. +/// Single allocation directly to the right size, no refcount overhead. +pub(crate) fn read_bytes_vec(cursor: &mut Cursor<&[u8]>) -> Result, MoonError> { + let len = read_u32(cursor)? as usize; + let pos = cursor.position() as usize; + let remaining = cursor.get_ref().len() - pos; + if len > remaining { + return Err(RdbError::Corrupted { + detail: format!( + "read_bytes_vec: length {} exceeds remaining {}", + len, remaining + ), + } + .into()); + } + let slice = &cursor.get_ref()[pos..pos + len]; + cursor.set_position((pos + len) as u64); + Ok(slice.to_vec()) +} + +/// Zero-copy read: returns a `Bytes` slice of the shared buffer (no heap alloc). +#[allow(dead_code)] +pub(crate) fn read_bytes_zero_copy( + cursor: &mut Cursor<&[u8]>, + shared_buf: &Bytes, +) -> Result { + let len = read_u32(cursor)? as usize; + let pos = cursor.position() as usize; + let remaining = cursor.get_ref().len() - pos; + if len > remaining { + return Err(RdbError::Corrupted { + detail: format!( + "read_bytes_zero_copy: length {} exceeds remaining {}", + len, remaining + ), + } + .into()); + } + cursor.set_position((pos + len) as u64); + Ok(shared_buf.slice(pos..pos + len)) } pub(crate) fn read_u32(cursor: &mut Cursor<&[u8]>) -> Result { diff --git a/src/server/conn/handler_monoio.rs b/src/server/conn/handler_monoio.rs index 220eb9a9c..f505fe18d 100644 --- a/src/server/conn/handler_monoio.rs +++ b/src/server/conn/handler_monoio.rs @@ -1174,28 +1174,6 @@ pub async fn handle_connection_sharded_monoio< continue; } - // --- BGSAVE: trigger per-shard cooperative snapshot --- - if cmd.eq_ignore_ascii_case(b"BGSAVE") { - let response = crate::command::persistence::bgsave_start_sharded( - &snapshot_trigger_tx, - num_shards, - ); - responses.push(response); - continue; - } - // SAVE -- not supported in sharded mode - if cmd.eq_ignore_ascii_case(b"SAVE") { - responses.push(Frame::Error(Bytes::from_static( - b"ERR SAVE not supported in sharded mode, use BGSAVE", - ))); - continue; - } - // LASTSAVE -- return timestamp of last successful save - if cmd.eq_ignore_ascii_case(b"LASTSAVE") { - responses.push(crate::command::persistence::handle_lastsave()); - continue; - } - // === ACL permission check (NOPERM gate) === // Exempt commands (AUTH, HELLO, QUIT, ACL) already handled above. { @@ -1239,6 +1217,39 @@ pub async fn handle_connection_sharded_monoio< } } + // --- BGSAVE: trigger per-shard cooperative snapshot --- + if cmd.eq_ignore_ascii_case(b"BGSAVE") { + let response = crate::command::persistence::bgsave_start_sharded( + &snapshot_trigger_tx, + num_shards, + ); + responses.push(response); + continue; + } + // SAVE -- not supported in sharded mode + if cmd.eq_ignore_ascii_case(b"SAVE") { + responses.push(Frame::Error(Bytes::from_static( + b"ERR SAVE not supported in sharded mode, use BGSAVE", + ))); + continue; + } + // LASTSAVE -- return timestamp of last successful save + if cmd.eq_ignore_ascii_case(b"LASTSAVE") { + responses.push(crate::command::persistence::handle_lastsave()); + continue; + } + if cmd.eq_ignore_ascii_case(b"BGREWRITEAOF") { + if let Some(ref tx) = aof_tx { + responses.push(crate::command::persistence::bgrewriteaof_start_sharded( + tx, + shard_databases.clone(), + )); + } else { + responses.push(Frame::Error(Bytes::from_static(b"ERR AOF is not enabled"))); + } + continue; + } + // --- MULTI --- if cmd.eq_ignore_ascii_case(b"MULTI") { if in_multi { diff --git a/src/server/conn/handler_sharded.rs b/src/server/conn/handler_sharded.rs index 404643f84..6be205cac 100644 --- a/src/server/conn/handler_sharded.rs +++ b/src/server/conn/handler_sharded.rs @@ -1214,6 +1214,13 @@ pub async fn handle_connection_sharded_inner< continue; } + // --- MULTI queue mode --- + if in_multi { + command_queue.push(frame); + responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); + continue; + } + // --- BGSAVE --- if cmd.eq_ignore_ascii_case(b"BGSAVE") { responses.push(crate::command::persistence::bgsave_start_sharded(&snapshot_trigger_tx, num_shards)); @@ -1227,11 +1234,17 @@ pub async fn handle_connection_sharded_inner< responses.push(crate::command::persistence::handle_lastsave()); continue; } - - // --- MULTI queue mode --- - if in_multi { - command_queue.push(frame); - responses.push(Frame::SimpleString(Bytes::from_static(b"QUEUED"))); + if cmd.eq_ignore_ascii_case(b"BGREWRITEAOF") { + if let Some(ref tx) = aof_tx { + responses.push(crate::command::persistence::bgrewriteaof_start_sharded( + tx, + shard_databases.clone(), + )); + } else { + responses.push(Frame::Error(Bytes::from_static( + b"ERR AOF is not enabled", + ))); + } continue; } diff --git a/src/server/listener.rs b/src/server/listener.rs index 561280259..0e856541d 100644 --- a/src/server/listener.rs +++ b/src/server/listener.rs @@ -411,7 +411,15 @@ pub async fn run_sharded( affinity_tracker: Arc>, ) -> anyhow::Result<()> { let addr = format!("{}:{}", config.bind, config.port); - let listener = monoio::net::TcpListener::bind(&addr)?; + // When per_shard_accept is true, each shard creates its own SO_REUSEPORT listener. + // We must NOT keep a central listener bound, because SO_REUSEPORT would distribute + // connections to it, but nobody would accept them (causing timeouts). + // Only bind the central listener when shards don't do per-shard accept. + let listener = if per_shard_accept { + None + } else { + Some(monoio::net::TcpListener::bind(&addr)?) + }; let num_shards = conn_txs.len(); info!("Listening on {} ({} shards, monoio)", addr, num_shards); @@ -447,12 +455,12 @@ pub async fn run_sharded( // If TLS listener is configured, select on both plain and TLS accepts if let Some(ref tls_listener) = tls_listener { monoio::select! { - // Plain TCP accept -- disabled when per-shard SO_REUSEPORT listeners are active. + // Plain TCP accept -- only active when central listener exists (non per-shard). result = async { - if per_shard_accept { - std::future::pending::>().await - } else { + if let Some(ref listener) = listener { listener.accept().await + } else { + std::future::pending::>().await } } => { match result { @@ -530,12 +538,12 @@ pub async fn run_sharded( } } else { monoio::select! { - // Plain TCP accept -- disabled when per-shard SO_REUSEPORT listeners are active. + // Plain TCP accept -- only active when central listener exists (non per-shard). result = async { - if per_shard_accept { - std::future::pending::>().await - } else { + if let Some(ref listener) = listener { listener.accept().await + } else { + std::future::pending::>().await } } => { match result { diff --git a/src/shard/dispatch.rs b/src/shard/dispatch.rs index 21dd2aadd..a6fbc24b6 100644 --- a/src/shard/dispatch.rs +++ b/src/shard/dispatch.rs @@ -370,6 +370,7 @@ mod tests { assert_eq!(key_to_shard(b"{tag}.key", 1), 0); } + #[cfg(feature = "runtime-tokio")] #[tokio::test] async fn test_pubsub_slot_waker() { let slot = Arc::new(PubSubResponseSlot::new(1)); @@ -387,6 +388,7 @@ mod tests { handle.await.unwrap(); } + #[cfg(feature = "runtime-tokio")] #[tokio::test] async fn test_pubsub_slot_multiple_shards() { let slot = Arc::new(PubSubResponseSlot::new(3)); @@ -412,6 +414,7 @@ mod tests { } } + #[cfg(feature = "runtime-tokio")] #[tokio::test] async fn test_pubsub_slot_already_ready() { // Slot with 0 pending should resolve immediately diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index bd7e3d844..5d2fa9c52 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -193,7 +193,7 @@ impl super::Shard { // Per-shard SO_REUSEPORT listener (Linux + monoio). // Each shard creates its own listener; the kernel distributes connections via SO_REUSEPORT. #[cfg(all(target_os = "linux", feature = "runtime-monoio"))] - let per_shard_monoio_listener: Option = { + let mut per_shard_monoio_listener: Option = { if let Some(ref addr) = bind_addr { match conn_accept::create_reuseport_socket(addr) { Ok(std_listener) => match monoio::net::TcpListener::from_std(std_listener) { @@ -335,6 +335,75 @@ impl super::Shard { #[cfg(feature = "runtime-monoio")] let pending_wakers: Rc>> = Rc::new(RefCell::new(Vec::new())); + // Spawn a dedicated accept loop for the per-shard monoio listener. + // This avoids the io_uring cancel/resubmit bug: monoio::select! drops and recreates + // the accept future each iteration (when periodic tick fires), causing in-flight + // io_uring ACCEPT operations to be cancelled asynchronously. Connections arriving + // during the cancel window are lost. A dedicated task keeps accept() alive + // continuously without cancellation. + #[cfg(all(target_os = "linux", feature = "runtime-monoio"))] + if let Some(listener) = per_shard_monoio_listener.take() { + let tls_cfg = tls_config.clone(); + let shard_dbs = shard_databases.clone(); + let dtx = dispatch_tx.clone(); + let ps = pubsub_arc.clone(); + let blk = blocking_rc.clone(); + let sd = shutdown.clone(); + let atx = aof_tx.clone(); + let trk = tracking_rc.clone(); + let lua = lua_rc.clone(); + let sc = script_cache_rc.clone(); + let acl = acl_table.clone(); + let rtcfg = runtime_config.clone(); + let svcfg = server_config.clone(); + let notifs = all_notifiers.to_vec(); + let snap_tx = snapshot_trigger_tx.clone(); + let rstate = repl_state.clone(); + let cstate = cluster_state.clone(); + let clock = cached_clock.clone(); + let rsm = remote_sub_map_arc.clone(); + let all_ps = all_pubsub_registries.to_vec(); + let all_rsm = all_remote_sub_maps.to_vec(); + let aff = affinity_tracker.clone(); + let pw = pending_wakers.clone(); + monoio::spawn(async move { + loop { + monoio::select! { + result = listener.accept() => { + match result { + Ok((stream, _addr)) => { + let std_stream = { + use std::os::unix::io::{IntoRawFd, FromRawFd}; + let fd = stream.into_raw_fd(); + // SAFETY: fd is valid, just transferred from monoio TcpStream + unsafe { std::net::TcpStream::from_raw_fd(fd) } + }; + conn_accept::spawn_monoio_connection( + std_stream, false, &tls_cfg, + &shard_dbs, &dtx, &ps, &blk, + &sd, &atx, &trk, &lua, &sc, + &acl, &rtcfg, &svcfg, ¬ifs, + &snap_tx, &rstate, &cstate, + &clock, &rsm, &all_ps, + &all_rsm, &aff, + shard_id, num_shards, config_port, + &pw, + ); + } + Err(e) => { + tracing::error!( + "Shard {}: per-shard accept error (monoio): {}", + shard_id, e + ); + } + } + } + _ = sd.cancelled() => break, + } + } + }); + } + loop { #[cfg(feature = "runtime-tokio")] tokio::select! { @@ -577,43 +646,11 @@ impl super::Shard { } } - // Monoio runtime: full event loop mirroring the tokio path. + // Monoio runtime: full event loop. + // Note: Per-shard SO_REUSEPORT accept runs in a dedicated spawned task (above) + // to avoid io_uring cancel/resubmit issues when select! drops accept futures. #[cfg(feature = "runtime-monoio")] monoio::select! { - // Per-shard SO_REUSEPORT accept (Linux only, monoio path) - result = async { - #[cfg(all(target_os = "linux", feature = "runtime-monoio"))] - if let Some(ref listener) = per_shard_monoio_listener { - return listener.accept().await; - } - // Never resolves on non-Linux or when per_shard_monoio_listener is None - std::future::pending::>().await - } => { - match result { - Ok((stream, _addr)) => { - // Convert monoio TcpStream -> std::net::TcpStream (same pattern as listener.rs) - let std_stream = { - use std::os::unix::io::{IntoRawFd, FromRawFd}; - let fd = stream.into_raw_fd(); - unsafe { std::net::TcpStream::from_raw_fd(fd) } - }; - conn_accept::spawn_monoio_connection( - std_stream, false, &tls_config, - &shard_databases, &dispatch_tx, &pubsub_arc, &blocking_rc, - &shutdown, &aof_tx, &tracking_rc, &lua_rc, &script_cache_rc, - &acl_table, &runtime_config, &server_config, &all_notifiers, - &snapshot_trigger_tx, &repl_state, &cluster_state, - &cached_clock, &remote_sub_map_arc, &all_pubsub_registries, - &all_remote_sub_maps, &affinity_tracker, - shard_id, num_shards, config_port, - &pending_wakers, - ); - } - Err(e) => { - tracing::error!("Shard {}: per-shard accept error (monoio): {}", shard_id, e); - } - } - } // Accept new connections from listener (MPSC fallback, always active on non-Linux) stream = conn_rx.recv_async() => { match stream { diff --git a/src/shard/mod.rs b/src/shard/mod.rs index b5a7dd472..33aeeddb1 100644 --- a/src/shard/mod.rs +++ b/src/shard/mod.rs @@ -58,8 +58,13 @@ impl Shard { /// /// Loads the per-shard RRDSHARD snapshot file first (if it exists), then replays /// the per-shard WAL for any commands written after the last snapshot. + /// + /// When `skip_wal` is true, per-shard WAL replay is skipped (used when + /// `appendonly=yes` because the global AOF is the source of truth for + /// write replay, and replaying per-shard WAL could cause double-replay). + /// /// Returns total keys loaded (snapshot + WAL replay). - pub fn restore_from_persistence(&mut self, persistence_dir: &str) -> usize { + pub fn restore_from_persistence(&mut self, persistence_dir: &str, skip_wal: bool) -> usize { use crate::persistence::snapshot::shard_snapshot_load; use crate::persistence::wal; @@ -80,9 +85,9 @@ impl Shard { } } - // Replay per-shard WAL + // Replay per-shard WAL (skip when global AOF is the write source of truth) let wal_file = wal::wal_path(dir, self.id); - if wal_file.exists() { + if !skip_wal && wal_file.exists() { match wal::replay_wal(&mut self.databases, &wal_file, &DispatchReplayEngine) { Ok(n) => { info!("Shard {}: replayed {} WAL commands", self.id, n); diff --git a/src/shard/shared_databases.rs b/src/shard/shared_databases.rs index 27be72726..801cea507 100644 --- a/src/shard/shared_databases.rs +++ b/src/shard/shared_databases.rs @@ -111,6 +111,14 @@ impl ShardDatabases { self.db_count } + /// Return a reference to all databases across all shards (for AOF rewrite). + /// + /// Callers iterate shards × dbs and acquire read locks individually. + #[inline] + pub fn all_shard_dbs(&self) -> &[Vec>] { + &self.shards + } + /// Collect snapshot metadata (segment counts, base timestamps) for a shard. /// /// Acquires brief read locks on each database to gather metadata needed diff --git a/src/storage/compact_value.rs b/src/storage/compact_value.rs index d39a88cf2..e830b3483 100644 --- a/src/storage/compact_value.rs +++ b/src/storage/compact_value.rs @@ -42,7 +42,7 @@ const HEAP_TAG_MASK: usize = 0x7; /// Thin wrapper for heap-allocated strings. /// At 24 bytes (Vec), this is smaller than RedisValue::String(Bytes) (~40 bytes) -/// and avoids the enum discriminant overhead. +/// and avoids the enum discriminant + refcount overhead. struct HeapString(Vec); /// Borrowed view of a CompactValue, for zero-copy read access. @@ -139,16 +139,13 @@ impl CompactValue { /// `RedisValue` enum wrapper (~40B savings per heap string). /// Collections are still stored as `Box`. pub fn from_redis_value(value: RedisValue) -> Self { - match &value { - RedisValue::String(s) if s.len() <= SSO_MAX_LEN => { - return Self::inline_string(s); - } - _ => {} - } - - // String heap path: store as Box<[u8]> directly (no RedisValue wrapper) - if let RedisValue::String(s) = &value { - return Self::heap_string(s); + // String fast path: inline SSO or zero-copy owned Bytes + if let RedisValue::String(s) = value { + return if s.len() <= SSO_MAX_LEN { + Self::inline_string(&s) + } else { + Self::heap_string_owned(s) + }; } // Collection heap path: store as Box @@ -183,10 +180,29 @@ impl CompactValue { } } - /// Create a heap-allocated string CompactValue from byte data. - /// Stores as `Box` — eliminates RedisValue enum wrapper. - /// HeapString is 24 bytes (Vec) vs RedisValue::String(Bytes) at ~40 bytes. + /// Create a heap-allocated string CompactValue from a byte slice (copies data). pub fn heap_string(data: &[u8]) -> Self { + Self::heap_string_vec(data.to_vec()) + } + + /// Create from owned Bytes (converts to Vec via Bytes::into for zero-copy + /// when Bytes has unique ownership, or copies when shared). + pub fn heap_string_owned(data: Bytes) -> Self { + // Bytes::into::> is zero-copy when refcount == 1, copies otherwise + Self::heap_string_vec(data.into()) + } + + /// Create from an owned Vec directly — no copy, no refcount. + /// This is the fastest path: one Box allocation for the HeapString wrapper. + /// Public for RDB loader fast path. + pub fn heap_string_vec_direct(data: Vec) -> Self { + if data.len() <= SSO_MAX_LEN { + return Self::inline_string(&data); + } + Self::heap_string_vec(data) + } + + fn heap_string_vec(data: Vec) -> Self { debug_assert!(data.len() > SSO_MAX_LEN); let str_len = data.len(); @@ -194,7 +210,7 @@ impl CompactValue { let copy_len = str_len.min(4); prefix[..copy_len].copy_from_slice(&data[..copy_len]); - let hs = Box::new(HeapString(data.to_vec())); + let hs = Box::new(HeapString(data)); let raw_ptr = Box::into_raw(hs) as usize; debug_assert!( raw_ptr & HEAP_TAG_MASK == 0, @@ -322,6 +338,7 @@ impl CompactValue { /// Get a mutable reference to the heap string bytes. /// Returns None for non-string types and inline values. + /// Note: returns `Option<&mut Vec>` for the underlying byte buffer. pub fn as_bytes_mut(&mut self) -> Option<&mut Vec> { if self.is_inline() { None diff --git a/src/storage/db.rs b/src/storage/db.rs index 541a740a4..63c8d143e 100644 --- a/src/storage/db.rs +++ b/src/storage/db.rs @@ -403,6 +403,37 @@ impl Database { self.data.insert(CompactKey::from(key), entry); } + /// Bulk-load insert: skip duplicate check, version tracking, and per-key memory accounting. + /// + /// Used exclusively during RDB/AOF restore where keys are guaranteed unique and + /// we recalculate `used_memory` once after the entire load completes. + /// This is ~3x faster than `set()` for large loads because it avoids: + /// - Hash lookup for existing key (can't exist during fresh load) + /// - `estimate_memory()` traversal per value + /// - Version increment logic + #[inline] + pub fn insert_for_load(&mut self, key: Bytes, entry: Entry) { + self.data.insert(CompactKey::from(key), entry); + } + + /// Recalculate `used_memory` by scanning all entries. Call once after bulk load. + pub fn recalculate_memory(&mut self) { + let mut total = 0usize; + for (key, entry) in self.data.iter() { + total += entry_overhead(key.as_bytes(), entry); + } + self.used_memory = total; + } + + /// Pre-size the internal hash table for an expected key count. + /// Eliminates segment splits during bulk load. + pub fn reserve(&mut self, additional: usize) { + if additional > self.data.len() { + let new_table = DashTable::with_capacity(additional); + self.data = new_table; + } + } + /// Remove a key and return its entry. No expiry check needed (DEL removes regardless). pub fn remove(&mut self, key: &[u8]) -> Option { if let Some(entry) = self.data.remove(key) { diff --git a/tests/integration.rs b/tests/integration.rs index e80212f1a..56b7ecb3b 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -2,6 +2,9 @@ //! //! Each test spawns a real TCP server on an OS-assigned port, connects with the //! `redis` crate client, exercises commands over real TCP, and shuts down cleanly. +//! +//! These tests require the tokio runtime (`run_with_shutdown` is tokio-only). +#![cfg(feature = "runtime-tokio")] use moon::runtime::cancel::CancellationToken; use moon::runtime::channel; diff --git a/tests/replication_test.rs b/tests/replication_test.rs index cd3e86699..e615d3d26 100644 --- a/tests/replication_test.rs +++ b/tests/replication_test.rs @@ -2,6 +2,9 @@ //! //! Tests REPLICAOF, REPLCONF, INFO replication, READONLY enforcement, //! and REPLICAOF NO ONE promotion -- using real TCP connections. +//! +//! These tests require the tokio runtime (`run_with_shutdown` is tokio-only). +#![cfg(feature = "runtime-tokio")] use moon::runtime::cancel::CancellationToken; use tokio::net::TcpListener;