From 07e555dc23e2c2ff209bb52f3ed7144a08471ae2 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Fri, 19 Dec 2025 18:34:45 +0800 Subject: [PATCH 1/6] fix: fix mixed_lru.py --- tests/mixed_lru.py | 188 +++++++++++++++++++++++++++++++-------------- 1 file changed, 132 insertions(+), 56 deletions(-) diff --git a/tests/mixed_lru.py b/tests/mixed_lru.py index bcdf1d8..9af5c8d 100644 --- a/tests/mixed_lru.py +++ b/tests/mixed_lru.py @@ -24,6 +24,12 @@ SOURCE_START, SOURCE_END = 0, NUM_PAGES // 2 DEST_START, DEST_END = NUM_PAGES // 2, NUM_PAGES + +def token_to_signature(token_val: int) -> int: + """将 token 映射到 uint8 的稳定签名值（用于数据完整性校验）。""" + # 产生尽量分散的 0..255 值，避免大量 token 映射到相同签名导致校验变弱 + return int(((token_val * 2654435761) ^ (token_val >> 16)) & 0xFF) + def test_mixed_lru_stability(): print("初始化测试环境...") @@ -66,20 +72,35 @@ def test_mixed_lru_stability(): print(f"初始化失败: {e}") return False + # LightMem 的基本语义是按 Block（pages_per_block 个 page）读写。 + # 测试侧也对齐到 Block 粒度，否则会出现：只写 1 页但触发 1 个 Block I/O， + # 再叠加并发与短时间窗口，容易偶现 read_success=0。 + pages_per_block = int(getattr(service, "_n", 0)) + if pages_per_block <= 0: + raise RuntimeError("PyLocalCacheService._n 不存在或无效，无法获取 pages_per_block") + + # 预留最后一个 source block 给 VIP，避免与写线程争抢同一 source 页 + vip_block_start = SOURCE_END - pages_per_block + usable_source_end = vip_block_start + if usable_source_end <= SOURCE_START: + raise RuntimeError("Source 区域不足以预留 VIP block") + # 4. 定义并发任务 stop_event = threading.Event() errors = [] - # 记录 Token -> Source Page Index 的映射 - # 用于验证读取回来的数据是否正确 + # 只记录“已确认落盘可 query”的 Token -> signature + # 读线程只从该集合抽样，避免把正在写/写失败(返回0)的 token 当作可读 token。 token_map = {} map_lock = threading.Lock() # 统计信息 stats = { - "written": 0, - "write_new": 0, # 新写入的 token - "write_update": 0, # 更新已有的 token + "written": 0, # 写入成功且可 query 的 token 次数 + "write_new": 0, # 新写入的 token + "write_update": 0, # 更新已有的 token + "write_retry": 0, # 因写入未落盘而重试次数 + "write_not_persisted": 0, # 写任务完成但 query 仍为 false 的次数 "read_success": 0, "read_miss_early": 0, # Query 阶段发现 Miss "read_miss_late": 0, # Read 阶段发现 Miss (并发淘汰) @@ -89,6 +110,9 @@ def test_mixed_lru_stability(): } stats_lock = threading.Lock() + # 写入就绪事件：避免 readers 在可读 token 为空时空转/偶现 0 成功读取 + ready_event = threading.Event() + # 写入线程 def writer_thread(tid): count = 0 @@ -97,11 +121,17 @@ def writer_thread(tid): base_token = tid * 1000000 token_pool_size = 100 # 每个线程 100 个不同的 token - # 每个线程独占一段 page 区域，避免多线程并发写同一个 page - # 将 SOURCE 区域平均分配给 4 个写线程 - num_writers = 4 - pages_per_writer = (SOURCE_END - SOURCE_START) // num_writers - thread_page_start = SOURCE_START + tid * pages_per_writer + # 每个线程分配若干 source block（对齐 pages_per_block），避免并发写同一 source block + source_block_starts = list(range(SOURCE_START, usable_source_end, pages_per_block)) + thread_blocks = source_block_starts[tid::4] + if not thread_blocks: + raise RuntimeError(f"Writer {tid} 没有可用的 source blocks") + + # 预构建 indexer tensor，避免循环中反复构造 + thread_indexers = [ + torch.arange(bs, bs + pages_per_block, dtype=torch.int32) + for bs in thread_blocks + ] while not stop_event.is_set(): try: @@ -111,38 +141,56 @@ def writer_thread(tid): data = [token_val] hash_128s = generate_cumulative_hashes(data) - # 映射到本线程独占的 page 区域 - page_idx = thread_page_start + (token_idx % pages_per_writer) - indexer = torch.tensor([page_idx], dtype=torch.int32) - # 写入前先更新内存中的数据为此token对应的特征值 - # 使用 token_val 的某个特征作为填充值,确保每个token有唯一的数据模式 - token_signature = (token_val // 1000) % 10 # 使用token的高位作为特征 - kvcache[page_idx].fill_(token_signature) + # 选择一个 source block 并填充整个 block 的数据（对齐底层 Block 粒度） + block_sel = token_idx % len(thread_blocks) + src_block_start = thread_blocks[block_sel] + indexer = thread_indexers[block_sel] + + token_signature = token_to_signature(token_val) + kvcache[src_block_start:src_block_start + pages_per_block].fill_(token_signature) - # 记录映射 (在写入前记录，虽然有微小的时间差，但只要不覆盖旧 Token 就行) + # 提交写任务并等待结束；注意：底层可能因为临时拥塞/写跳过导致 write 返回 0， + # 这种情况下 task 也会 ready，但 query 仍为 false。测试侧要识别并重试。 + max_retries = 10 + persisted = False is_new_token = False + for attempt in range(max_retries): + if attempt > 0: + with stats_lock: + stats["write_retry"] += 1 + time.sleep(0.01 * attempt) + + task = service.create(hash_128s=hash_128s, kv_page_indexer=indexer, mode="w") + while not task.ready(): + time.sleep(0.001) + + exists = service.query(hash_128s) + if exists and exists[0]: + persisted = True + break + + if not persisted: + with stats_lock: + stats["write_not_persisted"] += 1 + # 本轮写入没有真正落盘，不把 token 放进可读集合 + count += 1 + continue + with map_lock: is_new_token = (token_val not in token_map) - token_map[token_val] = (page_idx, token_signature) # 同时记录签名值 - # 限制 map 大小，移除太旧的（模拟应用层遗忘） - # 但为了测试 LRU，我们其实希望 map 里保留的比 disk 多，这样才能测出 miss + token_map[token_val] = token_signature + # 限制 map 大小（模拟应用层遗忘） if len(token_map) > 5000: - # 随机移除一些，或者移除最早的 - # 字典是插入有序的 (Python 3.7+) first_key = next(iter(token_map)) del token_map[first_key] - # 提交写任务 - task = service.create(hash_128s=hash_128s, kv_page_indexer=indexer, mode="w") - while not task.ready(): - time.sleep(0.001) - with stats_lock: stats["written"] += 1 - if is_new_token: - stats["write_new"] += 1 - else: - stats["write_update"] += 1 + stats["write_new"] += 1 if is_new_token else 0 + stats["write_update"] += 0 if is_new_token else 1 + + if stats["written"] >= 20: + ready_event.set() count += 1 if count % 100 == 0: @@ -162,6 +210,10 @@ def reader_thread(tid): while not stop_event.is_set(): try: + # 等待至少有一定数量的可读 token + if not ready_event.is_set(): + time.sleep(0.02) + continue # 随机选一个已知的 Token target_token = None expected_signature = -1 @@ -180,12 +232,12 @@ def reader_thread(tid): hot_tokens.remove(target_token) target_token = None else: - _, expected_signature = token_map[target_token] + expected_signature = token_map[target_token] if target_token is None: # 随机选一个 token，并可能将其加入热点列表 target_token = random.choice(list(token_map.keys())) - _, expected_signature = token_map[target_token] + expected_signature = token_map[target_token] # 10% 概率成为新的热点 if random.random() < 0.1 and len(hot_tokens) < 20: @@ -203,15 +255,29 @@ def reader_thread(tid): if not exists[0]: with stats_lock: stats["read_miss_early"] += 1 + # 已被淘汰，移出可读集合，避免反复命中 miss + with map_lock: + token_map.pop(target_token, None) + if target_token in hot_tokens: + try: + hot_tokens.remove(target_token) + except ValueError: + pass continue - # 尝试读取到 Dest 区域 - dest_page_idx = random.randint(DEST_START, DEST_END - 1) - indexer = torch.tensor([dest_page_idx], dtype=torch.int32) + # 尝试读取到 Dest 区域：对齐 Block 粒度，并为每个 reader 分配独占 dest blocks + dest_block_starts = list(range(DEST_START, DEST_END, pages_per_block)) + thread_dest_blocks = dest_block_starts[tid::4] + if not thread_dest_blocks: + raise RuntimeError(f"Reader {tid} 没有可用的 dest blocks") - # 先把 Dest Page 清零，防止残留数据干扰验证 - kvcache[dest_page_idx].zero_() - task = service.create(hash_128s=hash_128s, kv_page_indexer=indexer, mode="r") + # 选择一个 dest block（线程内轮转即可） + dest_block_start = thread_dest_blocks[count % len(thread_dest_blocks)] + dest_indexer = torch.arange(dest_block_start, dest_block_start + pages_per_block, dtype=torch.int32) + + # 先把 Dest Block 清零，防止残留数据干扰验证 + kvcache[dest_block_start:dest_block_start + pages_per_block].zero_() + task = service.create(hash_128s=hash_128s, kv_page_indexer=dest_indexer, mode="r") wait_start = time.time() while not task.ready(): @@ -229,17 +295,18 @@ def reader_thread(tid): # 验证数据 - 使用 token 的签名值 expected_val = expected_signature - # 使用 torch.all 进行严格的全量检查 - if torch.all(kvcache[dest_page_idx] == expected_val): + # 使用 torch.all 进行严格的全量检查（整个 block） + if torch.all(kvcache[dest_block_start:dest_block_start + pages_per_block] == expected_val): with stats_lock: stats["read_success"] += 1 else: with stats_lock: stats["data_mismatch"] += 1 print(f"[Reader {tid}] Data Mismatch! Token {target_token}, Expected Val {expected_val}") - # 打印实际值看看 (前10个和均值) - print(f" Actual mean: {torch.mean(kvcache[dest_page_idx].float()):.2f}") - print(f" First 10 vals: {kvcache[dest_page_idx][:10].tolist()}") + # 打印实际值看看 (第一页前10个和均值) + first_page = kvcache[dest_block_start] + print(f" Actual mean(first page): {torch.mean(first_page.float()):.2f}") + print(f" First 10 vals(first page): {first_page[:10].tolist()}") print(f" Task states: {task_states}") else: # 任务未完成、被中止或失败（可能是 partial read 导致的 abort，或并发淘汰） @@ -260,20 +327,24 @@ def reader_thread(tid): # VIP 保活线程 def vip_thread(): vip_token = 88888888 - # 使用最后一个page,避免与任何写线程冲突 - vip_page = SOURCE_END - 1 - vip_signature = (vip_token // 1000) % 10 + vip_signature = token_to_signature(vip_token) - # 先写入前更新内存数据 - kvcache[vip_page].fill_(vip_signature) + vip_indexer = torch.arange(vip_block_start, vip_block_start + pages_per_block, dtype=torch.int32) + kvcache[vip_block_start:vip_block_start + pages_per_block].fill_(vip_signature) print("[VIP] Writing VIP token...") vip_data = [vip_token] vip_hash_128s = generate_cumulative_hashes(vip_data) - task = service.create(hash_128s=vip_hash_128s, kv_page_indexer=torch.tensor([vip_page], dtype=torch.int32), mode="w") - while not task.ready(): - time.sleep(0.001) + # 写入 VIP 并确认落盘 + for attempt in range(10): + task = service.create(hash_128s=vip_hash_128s, kv_page_indexer=vip_indexer, mode="w") + while not task.ready(): + time.sleep(0.001) + exists = service.query(vip_hash_128s) + if exists and exists[0]: + break + time.sleep(0.02 * (attempt + 1)) # 循环保活 while not stop_event.is_set(): @@ -288,8 +359,8 @@ def vip_thread(): stats["vip_evicted"] += 1 # 如果被淘汰了，重新写入，继续测试 # print("[VIP] Evicted! Re-writing...") - kvcache[vip_page].fill_(vip_signature) # 重新填充数据 - task = service.create(hash_128s=vip_hash_128s, kv_page_indexer=torch.tensor([vip_page], dtype=torch.int32), mode="w") + kvcache[vip_block_start:vip_block_start + pages_per_block].fill_(vip_signature) + task = service.create(hash_128s=vip_hash_128s, kv_page_indexer=vip_indexer, mode="w") while not task.ready(): time.sleep(0.001) @@ -302,12 +373,17 @@ def vip_thread(): print("启动并发读写线程...") threads = [] - # 4 Writers + # 先启动 Writers，等积累一批可读 token 再启动 Readers（避免偶现 read_success=0） for i in range(4): t = threading.Thread(target=writer_thread, args=(i,)) threads.append(t) t.start() + # warmup：等待至少写入并确认落盘一定数量 token + warmup_timeout_s = 10 + print(f"Warmup: 等待至少 20 个 token 落盘（超时 {warmup_timeout_s}s）...") + ready_event.wait(timeout=warmup_timeout_s) + # 4 Readers for i in range(4): t = threading.Thread(target=reader_thread, args=(i,)) From eb264120719bd1f2396a8495919c610fc508fc38 Mon Sep 17 00:00:00 2001 From: blueswhen Date: Sat, 24 Jan 2026 21:26:07 +0800 Subject: [PATCH 2/6] multi node support --- .clang-format | 0 .clang-tidy | 0 .clangd | 0 .gitignore | 43 +- CMakeLists.txt | 3 + LICENSE | 0 README.md | 57 + pyproject.toml | 5 +- python/light_mem/__init__.py | 127 +- python/light_mem/etcd_coordinator.py | 556 ++++++ python/light_mem/etcd_v3_http.py | 378 +++++ python/light_mem/server_cli.py | 499 ++++++ src/CMakeLists.txt | 11 +- src/config.h | 0 src/core/cache_task.h | 12 + src/core/error.h | 0 src/core/task_queue.h | 0 src/pybind.cpp | 14 +- src/service/cache_service.h | 61 +- src/service/local_cache_service.h | 325 +++- src/storage/local_cache_index.cpp | 399 +++++ src/storage/local_cache_index.h | 95 ++ src/storage/local_storage_engine.h | 600 +++---- src/storage/local_storage_engine_files.cpp | 400 +++++ src/storage/local_storage_engine_journal.cpp | 334 ++++ src/storage/local_storage_engine_public.cpp | 615 +++++++ src/storage/local_storage_engine_recovery.cpp | 472 ++++++ src/storage/local_storage_engine_shard.cpp | 236 +++ src/storage/local_storage_wal.h | 62 + src/storage/redis_client.h | 545 ++++++ src/storage/storage_engine.h | 0 src/utils/fsync_compat.h | 28 + tests/mixed_lru.py | 0 tests/multi_node/harness.py | 741 ++++++++ tests/multi_node/run_all.py | 219 +++ tests/multi_node/stability_dynamic_nodes.py | 977 +++++++++++ .../multi_node/test_00_bootstrap_empty_dir.py | 206 +++ .../test_01_restart_recover_existing_dir.py | 312 ++++ ...st_02_assignment_balance_and_uniqueness.py | 244 +++ .../test_03_write_goes_to_owned_shards.py | 164 ++ .../test_04_cross_node_dedupe_same_hash.py | 247 +++ ...test_05_join_rebalance_minimal_movement.py | 277 +++ .../test_06_leave_lease_expire_reassign.py | 167 ++ ...est_07_redis_flush_and_recover_to_redis.py | 157 ++ .../test_08_cross_node_read_lru_window.py | 244 +++ .../test_09_crash_recovery_subprocess.py | 68 + .../test_10_watch_prefix_callback.py | 78 + tests/multi_node/test_11_crc_validation.py | 180 ++ tests/multi_node/test_12_redis_recovery.py | 861 ++++++++++ .../test_13_concurrent_init_shared_dir.py | 265 +++ tests/multi_node/worker_crash_node.py | 95 ++ tests/multi_node/worker_node_ops.py | 1494 +++++++++++++++++ tests/read.py | 1 + tests/test_utils.py | 0 tests/write.py | 1 + 55 files changed, 12363 insertions(+), 512 deletions(-) mode change 100644 => 100755 .clang-format mode change 100644 => 100755 .clang-tidy mode change 100644 => 100755 .clangd mode change 100644 => 100755 .gitignore mode change 100644 => 100755 CMakeLists.txt mode change 100644 => 100755 LICENSE mode change 100644 => 100755 README.md mode change 100644 => 100755 pyproject.toml mode change 100644 => 100755 python/light_mem/__init__.py create mode 100755 python/light_mem/etcd_coordinator.py create mode 100755 python/light_mem/etcd_v3_http.py create mode 100755 python/light_mem/server_cli.py mode change 100644 => 100755 src/CMakeLists.txt mode change 100644 => 100755 src/config.h mode change 100644 => 100755 src/core/cache_task.h mode change 100644 => 100755 src/core/error.h mode change 100644 => 100755 src/core/task_queue.h mode change 100644 => 100755 src/pybind.cpp mode change 100644 => 100755 src/service/cache_service.h mode change 100644 => 100755 src/service/local_cache_service.h create mode 100755 src/storage/local_cache_index.cpp create mode 100755 src/storage/local_cache_index.h mode change 100644 => 100755 src/storage/local_storage_engine.h create mode 100755 src/storage/local_storage_engine_files.cpp create mode 100755 src/storage/local_storage_engine_journal.cpp create mode 100755 src/storage/local_storage_engine_public.cpp create mode 100755 src/storage/local_storage_engine_recovery.cpp create mode 100755 src/storage/local_storage_engine_shard.cpp create mode 100755 src/storage/local_storage_wal.h create mode 100755 src/storage/redis_client.h mode change 100644 => 100755 src/storage/storage_engine.h create mode 100755 src/utils/fsync_compat.h mode change 100644 => 100755 tests/mixed_lru.py create mode 100755 tests/multi_node/harness.py create mode 100755 tests/multi_node/run_all.py create mode 100755 tests/multi_node/stability_dynamic_nodes.py create mode 100755 tests/multi_node/test_00_bootstrap_empty_dir.py create mode 100755 tests/multi_node/test_01_restart_recover_existing_dir.py create mode 100755 tests/multi_node/test_02_assignment_balance_and_uniqueness.py create mode 100755 tests/multi_node/test_03_write_goes_to_owned_shards.py create mode 100755 tests/multi_node/test_04_cross_node_dedupe_same_hash.py create mode 100755 tests/multi_node/test_05_join_rebalance_minimal_movement.py create mode 100755 tests/multi_node/test_06_leave_lease_expire_reassign.py create mode 100755 tests/multi_node/test_07_redis_flush_and_recover_to_redis.py create mode 100755 tests/multi_node/test_08_cross_node_read_lru_window.py create mode 100755 tests/multi_node/test_09_crash_recovery_subprocess.py create mode 100755 tests/multi_node/test_10_watch_prefix_callback.py create mode 100755 tests/multi_node/test_11_crc_validation.py create mode 100755 tests/multi_node/test_12_redis_recovery.py create mode 100755 tests/multi_node/test_13_concurrent_init_shared_dir.py create mode 100755 tests/multi_node/worker_crash_node.py create mode 100755 tests/multi_node/worker_node_ops.py mode change 100644 => 100755 tests/test_utils.py diff --git a/.clang-format b/.clang-format old mode 100644 new mode 100755 diff --git a/.clang-tidy b/.clang-tidy old mode 100644 new mode 100755 diff --git a/.clangd b/.clangd old mode 100644 new mode 100755 diff --git a/.gitignore b/.gitignore old mode 100644 new mode 100755 index ed8a765..cb2acc2 --- a/.gitignore +++ b/.gitignore @@ -1,21 +1,22 @@ -# Python-generated files -__pycache__/ -*.py[oc] -build/ -dist/ -wheels/ -*.egg-info - -# Virtual environments -.venv - -# IDE -/.cache -/.vscode - -# Generated -*.so -*.build -*.log -.py-build-cmake_cache/ -third_party/grpc/ +# Python-generated files +__pycache__/ +*.py[oc] +build/ +build_vscode/ +dist/ +wheels/ +*.egg-info + +# Virtual environments +.venv + +# IDE +/.cache +/.vscode + +# Generated +*.so +*.build +*.log +.py-build-cmake_cache/ +third_party/grpc/ diff --git a/CMakeLists.txt b/CMakeLists.txt old mode 100644 new mode 100755 index d4bec6d..28f9226 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,6 +1,9 @@ cmake_minimum_required(VERSION 3.25) project(light_mem LANGUAGES CXX) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + find_package(Threads REQUIRED) add_subdirectory(src) diff --git a/LICENSE b/LICENSE old mode 100644 new mode 100755 diff --git a/README.md b/README.md old mode 100644 new mode 100755 index b707878..aeee17c --- a/README.md +++ b/README.md @@ -98,6 +98,63 @@ Controls the maximum size of each cache block in megabytes (MB). - **Smaller blocks** (e.g., 16): More fine-grained control, better for random access, but higher overhead per operation - Must be set before starting the cache service + #### Index Persistence (Optional) + + LightMem can persist the hash index to an external index backend. When LightMem restarts (even after a process crash), as long as the backend is still running, LightMem can rebuild the in-memory hash index from it and continue using the existing local disk cache files. + + It also writes recovery metadata to each shard's `meta` file using **SuperBlock + Journal (truncate mode)**. If the index backend is not available (or index data is missing), LightMem can fall back to replaying the local journal to rebuild the index. + + Enable by passing parameters to `PyLocalCacheService`: + + - `index_endpoint` (optional; `""` disables index; `host` uses default ports; `host:port` is explicit) + - `coord_endpoints` (optional override; format: `host[:port][,host2[:port]...]`) + - `coord_ttl` / `coord_reconcile_sec` (optional; coordinator timing knobs) + + Single-host default ports (recommended): + + ```python + from light_mem import PyLocalCacheService + + svc = PyLocalCacheService( + kvcache_tensor=kvcache_tensor, + file=file, + index_endpoint="127.0.0.1", + ) + ``` + + Custom ports / explicit overrides: + + ```python + from light_mem import PyLocalCacheService + + svc = PyLocalCacheService( + kvcache_tensor=kvcache_tensor, + file=file, + index_endpoint="127.0.0.1:16379", + coord_endpoints="127.0.0.1:12379", + ) + ``` + + You can also use `lightmem_server` to start dependency services with one command. + +#### `lightmem_server` (one command) + +After installing LightMem, you can start dependency services via a single command: + +```bash +lightmem_server +``` + +It uses Docker Compose to run: +- an index backend for persistence / global dedupe +- a coordinator backend for multi-node shard coordination (optional) + +`lightmem_server` prints an example Python snippet you can use for LightMem clients (it uses `index_endpoint="127.0.0.1"` when using default ports). + +To enable multi-node shard coordination, pass `coord_endpoints` to `PyLocalCacheService`. + +Tip: you can also pass `--coord-ttl` / `--coord-reconcile-sec` to `lightmem_server` to include these values in the printed client snippet. + ## Quick Start ### Key Concepts diff --git a/pyproject.toml b/pyproject.toml old mode 100644 new mode 100755 index 9200fd6..2145f5c --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,9 @@ dependencies = [ 'numpy>=1.25.1' ] +[project.scripts] +lightmem_server = "light_mem.server_cli:main" + [build-system] requires = ["py-build-cmake~=0.4.3"] build-backend = "py_build_cmake.build" @@ -17,7 +20,7 @@ name = "light_mem" directory = "python" [tool.py-build-cmake.sdist] -include = ["CMakeLists.txt", "src"] +include = ["CMakeLists.txt", "src", "python"] [tool.py-build-cmake.cmake] minimum_version = "3.24" diff --git a/python/light_mem/__init__.py b/python/light_mem/__init__.py old mode 100644 new mode 100755 index 5064ff9..51bc264 --- a/python/light_mem/__init__.py +++ b/python/light_mem/__init__.py @@ -1,8 +1,9 @@ from enum import Enum -from typing import List +from typing import List, Optional import torch from . import light_mem +from .etcd_coordinator import maybe_start_etcd_coordinator class PyState(Enum): @@ -64,9 +65,39 @@ def page_already_list(self) -> List[int]: class PyLocalCacheService: """ 基于本地存储的异步数据存取服务 """ + @staticmethod + def _normalize_host_only(value: str) -> str: + v = (value or "").strip() + if not v: + return "" + if v.startswith("http://"): + v = v[len("http://"):] + elif v.startswith("https://"): + v = v[len("https://"):] + # If someone accidentally passes a URL, keep only the authority. + if "/" in v: + v = v.split("/", 1)[0] + + # Allow bracketed IPv6 (e.g. "[::1]"). + if ":" in v and not v.startswith("["): + raise ValueError("host must be host/ip without port") + return v + def __init__( - self, kvcache_tensor: torch.Tensor, file: str, storage_size: int = 32*1024*1024*1024, - num_shard: int = 32, num_worker: int = 16 + self, + kvcache_tensor: torch.Tensor, + file: str, + storage_size: int = 32*1024*1024*1024, + num_shard: int = 32, + num_worker: int = 16, + *, + index_endpoint: str = "", + index_prefix: str = "lightmem", + bandwidth_log: bool = True, + coord_endpoints: str = "", + coord_node_id: Optional[str] = None, + coord_ttl: int = 10, + coord_reconcile_sec: float = 10.0, ): """ 使用 PyLocalCacheService 来创建异步数据存取引擎 (基于本地磁盘) PyLocalCacheService 会直接从 kv cache 中读取数据，并异步地向 kv cache 中写入数据 @@ -80,6 +111,20 @@ def __init__( storage_size (int): 本地文件大小 (本地文件是分片存储的，这里表示总大小) kvcache_tensor (torch.Tensor): kvcache tensor，维度顺序为 [page, stride] num_worker (int): 工作线程数量 + index_endpoint (str): Redis 索引服务地址。 + - 若仅提供主机名/IP (如 "127.0.0.1")，默认使用 6379 端口，并尝试推断 coord_endpoints 为该主机:2379。 + - 若提供主机名:端口 (如 "127.0.0.1:6379")，则严格作为 Redis 地址，此时若需 ETCD 必须显式指定 coord_endpoints。 + - 为空则禁用索引和分布式功能。 + index_prefix (str): 索引和协调器使用的 Key 前缀。 + bandwidth_log (bool): 是否打印带宽统计日志。 + coord_endpoints (str): ETCD 协调服务地址 (用于分片所有权和分布式锁)。若为空且 index_endpoint 为纯主机名，会自动推断。 + coord_node_id (str, optional): 分布式协调中的当前节点 ID。 + - 若为 None，默认使用 socket.gethostname()。 + - 注意：如果在同一台机器上运行多个进程，必须手动指定不同的 ID 以避免冲突。 + coord_ttl (int): ETCD 租约 TTL (秒)，决定节点故障判定时间。默认 10 秒。 + coord_reconcile_sec (float): 协调器兜底轮询周期 (秒)。默认约为 TTL(秒)。 + - watch 可能因网络抖动/压缩等原因短暂不可用，此轮询用于保证最终收敛。 + - 若分片数量较大 (如 >500)，可适当调大以减轻 ETCD 压力。 """ if kvcache_tensor.dim() != 2: raise ValueError("kvcache_tensor 必须是二维张量，形如 [num of page, page size]") @@ -88,19 +133,87 @@ def __init__( num_pages_total = kvcache_tensor.shape[0] + index_endpoint_str = (index_endpoint or "").strip() + index_prefix_str = (index_prefix or "").strip() + coord_endpoints_str = (coord_endpoints or "").strip() + + # Parameter minimization: index_endpoint can be either: + # - "" (disable index) + # - "host" (treat as deps host; use default ports 6379/2379) + # - "host:port" (explicit index endpoint) + host_only = "" + if index_endpoint_str: + if index_endpoint_str.startswith("["): + # Bracketed IPv6 forms: + # - "[::1]" (host-only) + # - "[::1]:6379" (explicit) + if "]:" in index_endpoint_str: + host_only = "" + elif index_endpoint_str.endswith("]"): + host_only = self._normalize_host_only(index_endpoint_str) + else: + raise ValueError("index_endpoint IPv6 must be '[addr]' or '[addr]:port'") + elif ":" not in index_endpoint_str: + host_only = self._normalize_host_only(index_endpoint_str) + + if host_only: + index_endpoint_str = f"{host_only}:6379" + if not coord_endpoints_str: + coord_endpoints_str = f"{host_only}:2379" + self._c = light_mem.LocalCacheService( file=file, storage_size=storage_size, num_of_shard=num_shard, kvcache=kvcache_tensor, num_workers=num_worker, + index_endpoint=index_endpoint_str, + bandwidth_log=bool(bandwidth_log), + index_prefix=index_prefix_str, ) self._c.run() + + # Optional: start coordinator-driven shard ownership. + self._etcd_thread = maybe_start_etcd_coordinator( + self._c, + num_shards=num_shard, + endpoints=coord_endpoints_str, + index_prefix=index_prefix_str, + node_id=coord_node_id, + coord_ttl=int(coord_ttl), + coord_reconcile_sec=float(coord_reconcile_sec), + ) self._num_of_page_total: int = num_pages_total self._block_size: int = int(self._c.block_size()) self._page_size: int = int(self._c.page_size()) self._n: int = self._block_size // self._page_size + def close(self) -> None: + """Best-effort shutdown for background coordinator thread. + + Note: the underlying C++ LocalCacheService currently has no explicit stop() binding. + This method focuses on stopping the etcd lease keepalive thread so that keys can + expire and tests can deterministically emulate node leave. + """ + t = getattr(self, "_etcd_thread", None) + if t is None: + return + try: + t.stop() + except Exception: + pass + try: + t.join(timeout=2.0) + except Exception: + pass + self._etcd_thread = None + + def __del__(self): # pragma: no cover + try: + self.close() + except Exception: + pass + def _hash(self, hash_128s: List[int]) -> List[str]: """将 128 位哈希整数数组按照长度 n 进行划分，并生成块哈希。 @@ -163,6 +276,14 @@ def active_threads(self, mode: str) -> int: """统计当前处于执行中的读写任务数量，mode 使用 "r" 或 "w""" return int(self._c.active_create_count(mode)) + def eviction_count(self) -> int: + """返回本节点发生过的真实 LRU 淘汰次数（跨所有 shard 累加）。""" + return int(self._c.eviction_count()) + + def eviction_observed(self) -> bool: + """是否已经开始触发真实 LRU 淘汰（本节点）。""" + return bool(self._c.eviction_observed()) + def abort(self, t: PyTask): """ 终止一个任务的执行，此函数调用后， 工作线程不会立即停止工作(因为他们是异步的，可能有正在进行的任务)， diff --git a/python/light_mem/etcd_coordinator.py b/python/light_mem/etcd_coordinator.py new file mode 100755 index 0000000..8123a13 --- /dev/null +++ b/python/light_mem/etcd_coordinator.py @@ -0,0 +1,556 @@ +import threading +import time +import socket +import hashlib +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +from .etcd_v3_http import EtcdV3HttpClient + + +@dataclass(frozen=True) +class EtcdOptions: + endpoints: str + prefix: str + node_id: str + ttl: int + reconcile_interval_sec: float + + +def _parse_endpoints(endpoints: str) -> Tuple[str, int]: + # We pick the first endpoint. + first = (endpoints or "").split(",")[0].strip() + if not first: + return "127.0.0.1", 2379 + if ":" in first: + host, port = first.rsplit(":", 1) + try: + return host, int(port) + except Exception: + return host, 2379 + return first, 2379 + + +def _hrw_score(node_id: str, shard_id: int) -> int: + # Deterministic 64-bit score for HRW / Rendezvous hashing. + # + # Note: avoid simplistic rolling hashes here; for structured inputs like + # "node-{i}:{sid}" they can exhibit pathological ordering (one node wins + # almost all shards). We use stdlib blake2b to get a stable, well-distributed + # 64-bit score without external dependencies. + h = hashlib.blake2b(digest_size=8) + h.update(node_id.encode("utf-8")) + h.update(b":") + h.update(int(shard_id).to_bytes(4, byteorder="big", signed=False)) + return int.from_bytes(h.digest(), byteorder="big", signed=False) + + +class EtcdShardCoordinator(threading.Thread): + """Drive shard ownership via etcd and feed assignments into the C++ engine. + + This runs on Python side to avoid adding heavy C++ dependencies (gRPC/protobuf). + + Keys (under prefix): + - nodes/{node_id} (lease-bound) + - shards/{sid}/state (FREE|CLAIMED|DRAINING|SEALED) + - shards/{sid}/owner (node_id, lease-bound) + - shards/{sid}/handoff_to (node_id) # request old owner to drain/release + + Epoch fencing: + - epoch = create_revision(owner_key) after a successful claim + - passed down into C++ via update_shard_assignments() + - C++ persists epoch into shard superblock + writes an epoch marker to WAL + + Delayed handoff: + - old owner sets draining=1 locally when it sees handoff_to != self or desired owner changed + - coordinator waits until C++ reports inflight==0 before deleting owner key + """ + + def __init__(self, service, num_shards: int, opt: EtcdOptions): + super().__init__(daemon=True) + self._svc = service + self._num_shards = int(num_shards) + self._opt = opt + + self._prefix = (opt.prefix or "lightmem").rstrip("/") + + self._stop_evt = threading.Event() + + # Event-driven reconcile: etcd watch callbacks set this. + self._dirty_evt = threading.Event() + self._watch_ids: List[int] = [] + self._watch_enabled = False + self._watch_debounce_sec: float = 0.2 + + self._owned_epoch: Dict[int, int] = {} + self._draining: Dict[int, bool] = {} + + self._last_nodes: List[str] = [] + self._assignment: Dict[int, str] = {} + + + def stop(self) -> None: + self._stop_evt.set() + self._dirty_evt.set() + + def _client(self): + host, port = _parse_endpoints(self._opt.endpoints) + return EtcdV3HttpClient(host=host, port=port) + + def _k(self, suffix: str) -> str: + # Normalize to avoid accidental double slashes. + s = (suffix or "").lstrip("/") + return f"{self._prefix}/{s}".rstrip("/") + + def _nodes_prefix(self) -> str: + return f"{self._prefix}/nodes/" + + def _shards_prefix(self) -> str: + return f"{self._prefix}/shards/" + + def _ensure_state(self, client, sid: int) -> None: + state_key = self._k(f"shards/{sid}/state") + # Create default state FREE if absent + txn = client.transaction( + compare=[client.transactions.create(state_key) == 0], + success=[client.transactions.put(state_key, "FREE")], + failure=[], + ) + try: + txn + except Exception: + # Best-effort; if it fails, next reconcile will retry. + return + + def _get_nodes(self, client) -> List[str]: + prefix = self._nodes_prefix() + nodes: List[str] = [] + for _, meta in client.get_prefix(prefix): + try: + key = meta.key.decode("utf-8") + except Exception: + continue + if not key.startswith(prefix): + continue + node_id = key[len(prefix):] + if node_id: + nodes.append(node_id) + nodes.sort() + return nodes + + def _start_watchers(self, client) -> None: + # Best-effort: if watch APIs are unavailable, we fall back to polling. + self._watch_ids.clear() + self._watch_enabled = False + + add_watch = getattr(client, "add_watch_prefix_callback", None) + if not callable(add_watch): + return + + def _on_evt(*_args, **_kwargs) -> None: + # Coalesce bursts. + self._dirty_evt.set() + + try: + wid1 = add_watch(self._nodes_prefix(), _on_evt) + wid2 = add_watch(self._shards_prefix(), _on_evt) + for wid in (wid1, wid2): + if isinstance(wid, int): + self._watch_ids.append(wid) + self._watch_enabled = True + except Exception: + self._watch_ids.clear() + self._watch_enabled = False + + def _stop_watchers(self, client) -> None: + cancel = getattr(client, "cancel_watch", None) + if callable(cancel): + for wid in list(self._watch_ids): + try: + cancel(wid) + except Exception: + pass + self._watch_ids.clear() + self._watch_enabled = False + + def _reconcile_once(self, client, lease) -> None: + nodes = self._get_nodes(client) + + # Recompute assignment only when node membership changes. + self._recompute_assignment_on_membership_change(client, nodes) + + shard_ids: List[int] = [] + epochs: List[int] = [] + draining: List[int] = [] + + newly_claimed: List[int] = [] + + for sid in range(self._num_shards): + desired = self._assignment.get(sid) or self._desired_owner(nodes, sid) + + state_key = self._k(f"shards/{sid}/state") + owner_key = self._k(f"shards/{sid}/owner") + handoff_key = self._k(f"shards/{sid}/handoff_to") + + state = self._get_text(client, state_key) + if state is None: + # Lazy init to reduce etcd traffic. + self._ensure_state(client, sid) + state = self._get_text(client, state_key) + state = state or "FREE" + + if state == "SEALED": + # never writable + self._owned_epoch.pop(sid, None) + self._draining.pop(sid, None) + continue + + owner = self._get_text(client, owner_key) + handoff_to = self._get_text(client, handoff_key) + + if desired == self._opt.node_id: + if owner is None: + ep = self._claim(client, sid, lease) + if ep: + if sid not in self._owned_epoch: + newly_claimed.append(sid) + self._owned_epoch[sid] = ep + self._draining[sid] = False + elif owner != self._opt.node_id: + # Ask current owner to drain/release. + self._request_handoff(client, sid) + else: + # We own it. + v, meta = client.get(owner_key) + # IMPORTANT: use create_revision as epoch fencing. + # mod_revision can change on lease refresh (PUT with same value), which would + # otherwise cause epoch churn and heavy metadata write contention in C++. + ep = int(getattr(meta, "create_revision", 0) or 0) + if not ep: + ep = int(getattr(meta, "mod_revision", 0) or 0) + if ep: + self._owned_epoch[sid] = ep + # If someone else requested handoff, start draining. + # If this node is the desired owner, clear stale handoff_to to avoid + # permanently-drained shards after transient membership changes. + if handoff_to and handoff_to != self._opt.node_id: + if desired == self._opt.node_id: + try: + client.delete(handoff_key) + except Exception: + pass + self._draining[sid] = False + else: + self._draining[sid] = True + else: + self._draining[sid] = False + else: + if owner == self._opt.node_id: + # We should no longer own it; start draining and release when safe. + self._draining[sid] = True + self._release_if_safe(client, sid, desired) + else: + self._owned_epoch.pop(sid, None) + self._draining.pop(sid, None) + + # Push assignments into C++. + for sid, ep in self._owned_epoch.items(): + shard_ids.append(int(sid)) + epochs.append(int(ep)) + draining.append(1 if self._draining.get(sid, False) else 0) + + self._svc.update_shard_assignments(shard_ids, epochs, draining) + + # On newly claimed shards, recover Redis index with a smart policy. + for sid in newly_claimed: + try: + self._svc.recover_shard_to_redis_smart(int(sid)) + except Exception: + pass + + def _desired_owner(self, nodes: List[str], sid: int) -> Optional[str]: + if not nodes: + return None + best_node = None + best_score = None + for n in nodes: + sc = _hrw_score(n, sid) + if best_score is None or sc > best_score: + best_score = sc + best_node = n + return best_node + + def _recompute_assignment_on_membership_change(self, client, nodes: List[str]) -> None: + if nodes == self._last_nodes and self._assignment: + return + + self._last_nodes = list(nodes) + self._assignment.clear() + if not nodes: + return + + # HRW / Rendezvous hashing assignment: + # - Deterministic mapping shard -> node based solely on (node_id, shard_id) + # - Minimal movement when nodes join/leave + # - Stable when membership unchanged + for sid in range(self._num_shards): + self._assignment[sid] = self._desired_owner(nodes, sid) # type: ignore[assignment] + + def _get_text(self, client, key: str) -> Optional[str]: + v, _ = client.get(key) + if v is None: + return None + try: + return v.decode("utf-8") + except Exception: + return None + + def _claim(self, client, sid: int, lease) -> Optional[int]: + owner_key = self._k(f"shards/{sid}/owner") + state_key = self._k(f"shards/{sid}/state") + handoff_key = self._k(f"shards/{sid}/handoff_to") + + # Only claim if no owner and not SEALED. + txn_ok, _ = client.transaction( + compare=[ + client.transactions.create(owner_key) == 0, + client.transactions.value(state_key) != b"SEALED", + ], + success=[ + client.transactions.put(owner_key, self._opt.node_id, lease=lease), + client.transactions.put(state_key, "CLAIMED"), + client.transactions.delete(handoff_key), + ], + failure=[], + ) + if not txn_ok: + return None + + # epoch = create_revision(owner_key) (stable across lease refresh) + v, meta = client.get(owner_key) + if v is None or meta is None: + return None + ep = int(getattr(meta, "create_revision", 0) or 0) + if not ep: + ep = int(getattr(meta, "mod_revision", 0) or 0) + return ep + + def _request_handoff(self, client, sid: int) -> None: + handoff_key = self._k(f"shards/{sid}/handoff_to") + # Always overwrite to avoid stale handoff targets pinning shards in draining state. + try: + client.put(handoff_key, self._opt.node_id) + except Exception: + # Best-effort; next reconcile will retry. + pass + + def _release_if_safe(self, client, sid: int, desired_owner: Optional[str]) -> None: + owner_key = self._k(f"shards/{sid}/owner") + state_key = self._k(f"shards/{sid}/state") + + # Only release if no inflight writes. + inflight = int(self._svc.shard_inflight(int(sid))) + if inflight != 0: + return + + # Best-effort state update + owner delete. + try: + if desired_owner and desired_owner != self._opt.node_id: + client.put(state_key, "FREE") + client.delete(owner_key) + self._owned_epoch.pop(sid, None) + self._draining.pop(sid, None) + except Exception: + return + + def run(self) -> None: + def _best_effort_cleanup(c) -> None: + # Best-effort cleanup to reduce stale membership/ownership during + # process restarts (e.g., pipeline phases). This is intentionally + # conservative: only delete owner keys if they still point to us. + try: + node_key_local = self._k(f"nodes/{self._opt.node_id}") + try: + c.delete(node_key_local) + except Exception: + pass + + for sid in range(int(self._num_shards)): + owner_key = self._k(f"shards/{int(sid)}/owner") + state_key = self._k(f"shards/{int(sid)}/state") + try: + v, _m = c.get(owner_key) + owner = None + if v is not None: + try: + owner = v.decode("utf-8") + except Exception: + owner = None + if owner == self._opt.node_id: + try: + c.put(state_key, "FREE") + except Exception: + pass + try: + c.delete(owner_key) + except Exception: + pass + except Exception: + continue + except Exception: + return + + client = self._client() + lease = client.lease(self._opt.ttl) + + # Register node key bound to the lease (one-time); keepalive refreshes the lease. + node_key = self._k(f"nodes/{self._opt.node_id}") + client.put(node_key, "1", lease=lease) + + self._start_watchers(client) + # Force an initial reconcile. + self._dirty_evt.set() + + keepalive_interval = max(1.0, float(self._opt.ttl) / 3.0) + fallback_reconcile_interval = max(1.0, float(self._opt.reconcile_interval_sec)) + + next_keepalive = 0.0 + next_fallback_reconcile = 0.0 + + try: + while not self._stop_evt.is_set(): + now = time.time() + + # Keep the lease alive. + if now >= next_keepalive: + try: + # Use one-shot keepalive against the grpc-gateway streaming endpoint. + # This avoids periodic re-PUTs of all shard owner keys, which would + # otherwise generate watch events and amplify reconcile load. + lease = client.lease_keepalive_once(lease) + except Exception: + # Re-establish client/lease/watch on failures. + try: + self._stop_watchers(client) + except Exception: + pass + client = self._client() + lease = client.lease(self._opt.ttl) + client.put(node_key, "1", lease=lease) + # Re-attach owned shard owner keys to the new lease. + for sid in list(self._owned_epoch.keys()): + owner_key = self._k(f"shards/{int(sid)}/owner") + try: + client.put(owner_key, self._opt.node_id, lease=lease) + except Exception: + pass + self._start_watchers(client) + self._dirty_evt.set() + next_keepalive = now + keepalive_interval + + # Decide how long to wait for events. + timeout = max(0.0, min(next_keepalive - now, next_fallback_reconcile - now)) + if self._dirty_evt.is_set(): + timeout = 0.0 + + fired = self._dirty_evt.wait(timeout=timeout) + if self._stop_evt.is_set(): + break + + # Debounce watch bursts. + if fired: + self._dirty_evt.clear() + time.sleep(self._watch_debounce_sec) + + now = time.time() + should_reconcile = fired or (now >= next_fallback_reconcile) + if not should_reconcile: + continue + + try: + self._reconcile_once(client, lease) + except Exception: + # Fail-closed on coordination errors: stop issuing new writes until we can reconcile again. + try: + self._svc.update_shard_assignments([], [], []) + except Exception: + pass + + next_fallback_reconcile = time.time() + fallback_reconcile_interval + finally: + try: + self._stop_watchers(client) + except Exception: + pass + _best_effort_cleanup(client) + +# Cluster-level guard: ensure all nodes agree on shard ID space size. +# +# Without this, different nodes may operate on different shard id ranges, +# leading to inconsistent ownership keys and unsafe writes. +def _ensure_cluster_num_shards(client, *, prefix: str, expected: int) -> None: + key = f"{prefix.rstrip('/')}/config/num_shards" + + # First node wins by creating the key. Others validate it matches. + try: + ok, _ = client.transaction( + compare=[client.transactions.create(key) == 0], + success=[client.transactions.put(key, str(int(expected)))], + failure=[], + ) + except Exception as e: + raise RuntimeError(f"etcd transaction failed while checking {key}: {e}") from e + + if ok: + return + + v, _ = client.get(key) + if v is None: + # Extremely unlikely (key existed for compare but missing now); treat as misconfiguration. + raise RuntimeError(f"cluster shard config key disappeared: {key}") + try: + actual = int(v.decode("utf-8").strip()) + except Exception: + raise RuntimeError(f"cluster shard config key is not an int: {key}={v!r}") + + if int(actual) != int(expected): + raise RuntimeError( + f"num_shards mismatch across nodes: expected {int(expected)} but etcd has {int(actual)} at {key}" + ) + +def maybe_start_etcd_coordinator( + service, + num_shards: int, + *, + endpoints: str, + index_prefix: str, + node_id: Optional[str], + coord_ttl: int, + coord_reconcile_sec: float, +) -> Optional[EtcdShardCoordinator]: + endpoints = (endpoints or "").strip() + if not endpoints: + return None + + node_id_final = (node_id or "").strip() or (socket.gethostname() or "unknown").strip() + prefix = (index_prefix or "").strip() + if not prefix: + raise ValueError("index_prefix must be non-empty when etcd coordination is enabled") + + ttl = int(coord_ttl) + interval = float(coord_reconcile_sec) + + host, port = _parse_endpoints(endpoints) + client = EtcdV3HttpClient(host=host, port=port) + _ensure_cluster_num_shards(client, prefix=prefix, expected=int(num_shards)) + + opt = EtcdOptions( + endpoints=endpoints, + prefix=prefix, + node_id=node_id_final, + ttl=int(ttl), + reconcile_interval_sec=float(interval), + ) + t = EtcdShardCoordinator(service=service, num_shards=int(num_shards), opt=opt) + t.start() + return t diff --git a/python/light_mem/etcd_v3_http.py b/python/light_mem/etcd_v3_http.py new file mode 100755 index 0000000..8766535 --- /dev/null +++ b/python/light_mem/etcd_v3_http.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +import base64 +import json +import socket +import threading +import time +import urllib.error +import urllib.request +from dataclasses import dataclass +from typing import Iterable, Iterator, Optional, Tuple, Dict, Callable + + +def _b64e(b: bytes) -> str: + return base64.b64encode(b).decode("ascii") + + +def _b64d(s: str) -> bytes: + return base64.b64decode(s.encode("ascii")) + + +def _prefix_range_end(prefix: bytes) -> bytes: + # Standard etcd prefix end calculation: increment the last non-0xFF byte. + if not prefix: + return b"\0" + ba = bytearray(prefix) + for i in range(len(ba) - 1, -1, -1): + if ba[i] != 0xFF: + ba[i] += 1 + return bytes(ba[: i + 1]) + return b"\0" + + +@dataclass(frozen=True) +class EtcdMeta: + key: bytes + create_revision: int + mod_revision: int + version: int + + +class _CreateCompareBuilder: + def __init__(self, key: str): + self._key = key + + def __eq__(self, other: object): + # create(key) == 0 -> VERSION(key) == 0 + try: + v = int(other) # type: ignore[arg-type] + except Exception: + v = 0 + return _TxnCompare.version_equal(self._key, v) + + +class _ValueCompareBuilder: + def __init__(self, key: str): + self._key = key + + def __eq__(self, other: object): + b = other if isinstance(other, (bytes, bytearray)) else str(other).encode("utf-8") + return _TxnCompare.value_equal(self._key, bytes(b)) + + def __ne__(self, other: object): + b = other if isinstance(other, (bytes, bytearray)) else str(other).encode("utf-8") + return _TxnCompare.value_not_equal(self._key, bytes(b)) + + +@dataclass(frozen=True) +class _TxnCompare: + target: str + key: str + result: str + version: Optional[int] = None + value: Optional[bytes] = None + + @staticmethod + def version_equal(key: str, version: int) -> "_TxnCompare": + return _TxnCompare(target="VERSION", key=key, result="EQUAL", version=int(version)) + + @staticmethod + def value_equal(key: str, value: bytes) -> "_TxnCompare": + return _TxnCompare(target="VALUE", key=key, result="EQUAL", value=value) + + @staticmethod + def value_not_equal(key: str, value: bytes) -> "_TxnCompare": + return _TxnCompare(target="VALUE", key=key, result="NOT_EQUAL", value=value) + + +@dataclass(frozen=True) +class _TxnPut: + key: str + value: str + lease_id: int = 0 + + +@dataclass(frozen=True) +class _TxnDelete: + key: str + + +class _TxnOps: + def create(self, key: str) -> _CreateCompareBuilder: + return _CreateCompareBuilder(key) + + def value(self, key: str) -> _ValueCompareBuilder: + return _ValueCompareBuilder(key) + + def put(self, key: str, value: str, *, lease=None) -> _TxnPut: + lease_id = 0 + if lease is not None: + lease_id = int(getattr(lease, "id", lease) or 0) + return _TxnPut(key=key, value=str(value), lease_id=lease_id) + + def delete(self, key: str) -> _TxnDelete: + return _TxnDelete(key=key) + + +@dataclass +class EtcdLease: + id: int + ttl: int + + +class _WatchThread(threading.Thread): + def __init__(self, host: str, port: int, key_b64: str, range_end_b64: str, callback: Callable): + super().__init__(daemon=True) + self._url = f"http://{host}:{port}/v3/watch" + self._payload = json.dumps({ + "create_request": { + "key": key_b64, + "range_end": range_end_b64, + } + }).encode("utf-8") + self._cb = callback + self._stop_evt = threading.Event() + + def stop(self): + self._stop_evt.set() + + def run(self): + while not self._stop_evt.is_set(): + try: + req = urllib.request.Request( + self._url, + data=self._payload, + headers={"Content-Type": "application/json"}, + method="POST", + ) + # No timeout for long streaming connection + with urllib.request.urlopen(req, timeout=None) as resp: + for line in resp: + if self._stop_evt.is_set(): + break + line = line.strip() + if not line: + continue + try: + msg = json.loads(line.decode("utf-8")) + # The 'result' field contains 'events' if changed + res = msg.get("result", {}) + if "events" in res: + self._cb(res["events"]) + except Exception: + # Ignore malformed chunks + pass + except Exception: + # Connection lost or failed, backoff and retry + if self._stop_evt.is_set(): + break + time.sleep(1.0) + + +class EtcdV3HttpClient: + """Minimal etcd v3 client via HTTP/JSON gateway. + + This avoids protobuf/gRPC client dependencies, and works with etcd 3.x that exposes + the grpc-gateway endpoints (e.g. /v3/kv/range, /v3/kv/txn, /v3/lease/grant). + """ + + def __init__(self, *, host: str = "127.0.0.1", port: int = 2379, timeout_s: float = 3.0): + self._host = host + self._port = int(port) + self._timeout_s = float(timeout_s) + + self.transactions = _TxnOps() + + self._watch_lock = threading.Lock() + self._watch_counter = 0 + self._active_watches: Dict[int, _WatchThread] = {} + + def _url(self, path: str) -> str: + p = "/" + (path or "").lstrip("/") + return f"http://{self._host}:{self._port}{p}" + + def _post_json(self, path: str, payload: dict) -> dict: + data = json.dumps(payload).encode("utf-8") + req = urllib.request.Request( + self._url(path), + data=data, + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=self._timeout_s) as resp: + raw = resp.read() + except (urllib.error.URLError, socket.timeout) as e: + raise ConnectionError(f"etcd http request failed: {e}") from e + try: + return json.loads(raw.decode("utf-8")) if raw else {} + except Exception as e: + raise RuntimeError(f"etcd http response is not json: {raw[:200]!r}") from e + + def lease(self, ttl: int) -> EtcdLease: + out = self._post_json("/v3/lease/grant", {"TTL": int(ttl)}) + lease_id = int(out.get("ID", 0) or 0) + if lease_id == 0: + raise RuntimeError(f"failed to grant lease: {out}") + return EtcdLease(id=lease_id, ttl=int(ttl)) + + def lease_keepalive_once(self, lease: EtcdLease | int, *, timeout_s: Optional[float] = None) -> EtcdLease: + """Best-effort lease keepalive via grpc-gateway streaming endpoint. + + etcd exposes LeaseKeepAlive as a streaming RPC; the HTTP gateway returns a + stream of JSON objects. To keep this client dependency-free, we do a + one-shot keepalive: open a request with a single keepalive message, + read one response frame, then close. + + Returns an updated EtcdLease with refreshed TTL on success. + """ + + lease_id = int(getattr(lease, "id", lease) or 0) + if lease_id <= 0: + raise ValueError(f"invalid lease id: {lease_id}") + + data = json.dumps({"ID": str(int(lease_id))}).encode("utf-8") + req = urllib.request.Request( + self._url("/v3/lease/keepalive"), + data=data, + headers={"Content-Type": "application/json"}, + method="POST", + ) + + effective_timeout = self._timeout_s if timeout_s is None else float(timeout_s) + try: + with urllib.request.urlopen(req, timeout=effective_timeout) as resp: + # Read a single streaming frame. + raw_line = resp.readline() + except (urllib.error.URLError, socket.timeout) as e: + raise ConnectionError(f"etcd lease keepalive failed: {e}") from e + + if not raw_line: + raise RuntimeError("etcd lease keepalive returned empty response") + + try: + msg = json.loads(raw_line.decode("utf-8")) + except Exception as e: + raise RuntimeError(f"etcd lease keepalive response is not json: {raw_line[:200]!r}") from e + + # grpc-gateway stream objects typically look like {"result": {...}, "error": {...}} + res = msg.get("result", msg) + try: + ttl = int(res.get("TTL", 0) or 0) + except Exception: + ttl = 0 + if ttl <= 0: + raise RuntimeError(f"etcd lease keepalive returned invalid TTL: {msg}") + return EtcdLease(id=int(lease_id), ttl=int(ttl)) + + def get(self, key: str) -> Tuple[Optional[bytes], Optional[EtcdMeta]]: + k = key.encode("utf-8") + out = self._post_json("/v3/kv/range", {"key": _b64e(k), "limit": 1}) + kvs = out.get("kvs") or [] + if not kvs: + return None, None + kv = kvs[0] + try: + meta = EtcdMeta( + key=_b64d(kv.get("key", "")), + create_revision=int(kv.get("create_revision", 0) or 0), + mod_revision=int(kv.get("mod_revision", 0) or 0), + version=int(kv.get("version", 0) or 0), + ) + v = _b64d(kv.get("value", "")) if kv.get("value") is not None else None + except Exception: + return None, None + return v, meta + + def get_prefix(self, prefix: str) -> Iterator[Tuple[Optional[bytes], EtcdMeta]]: + p = prefix.encode("utf-8") + range_end = _prefix_range_end(p) + out = self._post_json( + "/v3/kv/range", + { + "key": _b64e(p), + "range_end": _b64e(range_end), + }, + ) + for kv in out.get("kvs") or []: + try: + meta = EtcdMeta( + key=_b64d(kv.get("key", "")), + create_revision=int(kv.get("create_revision", 0) or 0), + mod_revision=int(kv.get("mod_revision", 0) or 0), + version=int(kv.get("version", 0) or 0), + ) + v = _b64d(kv.get("value", "")) if kv.get("value") is not None else None + except Exception: + continue + yield v, meta + + def put(self, key: str, value: str, *, lease=None) -> None: + lease_id = 0 + if lease is not None: + lease_id = int(getattr(lease, "id", lease) or 0) + payload = { + "key": _b64e(key.encode("utf-8")), + "value": _b64e(str(value).encode("utf-8")), + } + if lease_id: + payload["lease"] = lease_id + _ = self._post_json("/v3/kv/put", payload) + + def delete(self, key: str) -> None: + payload = {"key": _b64e(key.encode("utf-8"))} + _ = self._post_json("/v3/kv/deleterange", payload) + + def transaction(self, *, compare: Iterable[_TxnCompare], success: Iterable[object], failure: Iterable[object]): + def _cmp(c: _TxnCompare) -> dict: + d: dict = {"target": c.target, "key": _b64e(c.key.encode("utf-8")), "result": c.result} + if c.version is not None: + d["version"] = int(c.version) + if c.value is not None: + d["value"] = _b64e(bytes(c.value)) + return d + + def _req(op: object) -> dict: + if isinstance(op, _TxnPut): + r = { + "request_put": { + "key": _b64e(op.key.encode("utf-8")), + "value": _b64e(op.value.encode("utf-8")), + } + } + if op.lease_id: + r["request_put"]["lease"] = int(op.lease_id) + return r + if isinstance(op, _TxnDelete): + return {"request_delete_range": {"key": _b64e(op.key.encode("utf-8"))}} + raise TypeError(f"unsupported txn op: {type(op)!r}") + + payload = { + "compare": [_cmp(c) for c in compare], + "success": [_req(o) for o in success], + "failure": [_req(o) for o in failure], + } + out = self._post_json("/v3/kv/txn", payload) + return bool(out.get("succeeded", False)), out.get("responses") or [] + + def add_watch_prefix_callback(self, prefix: str, callback: Callable) -> int: + p = prefix.encode("utf-8") + range_end = _prefix_range_end(p) + key_b64 = _b64e(p) + end_b64 = _b64e(range_end) + + t = _WatchThread(self._host, self._port, key_b64, end_b64, callback) + t.start() + + with self._watch_lock: + self._watch_counter += 1 + wid = self._watch_counter + self._active_watches[wid] = t + return wid + + def cancel_watch(self, watch_id: int) -> None: + with self._watch_lock: + t = self._active_watches.pop(watch_id, None) + if t: + t.stop() diff --git a/python/light_mem/server_cli.py b/python/light_mem/server_cli.py new file mode 100755 index 0000000..c65665f --- /dev/null +++ b/python/light_mem/server_cli.py @@ -0,0 +1,499 @@ +import argparse +import os +import shutil +import subprocess +import sys +import tempfile +import signal +import time +from textwrap import dedent + + +def _which(cmd: str) -> str | None: + return shutil.which(cmd) + + +def _run(cmd: list[str]) -> int: + try: + p = subprocess.run(cmd, check=False) + return int(p.returncode) + except FileNotFoundError: + return 127 + + +def _docker_compose_cmd() -> list[str] | None: + # Prefer: docker compose + if _which("docker"): + # Best-effort: if `docker compose version` works, use it. + rc = _run(["docker", "compose", "version"]) + if rc == 0: + return ["docker", "compose"] + + # Fallback: legacy docker-compose + if _which("docker-compose"): + return ["docker-compose"] + + return None + + +def _stop_docker_services(*, purge_volumes: bool) -> int: + docker = _which("docker") + if not docker: + sys.stderr.write("docker is required for lightmem_server --stop.\n") + return 2 + + # Containers are created with fixed names. + containers = ["lightmem-index", "lightmem-coord"] + rm = subprocess.run([docker, "rm", "-f", *containers], capture_output=True, text=True, check=False) + + # If containers do not exist, treat as success. + if rm.returncode != 0: + err = (rm.stderr or "") + (rm.stdout or "") + lowered = err.lower() + if "no such container" not in lowered: + sys.stderr.write(err) + sys.stderr.write("Failed to stop containers.\n") + return int(rm.returncode) + + if purge_volumes: + # Volumes are named explicitly by our compose YAML. + vols = ["lightmem-redis-data", "lightmem-etcd-data"] + vrm = subprocess.run([docker, "volume", "rm", "-f", *vols], capture_output=True, text=True, check=False) + if vrm.returncode != 0: + err = (vrm.stderr or "") + (vrm.stdout or "") + lowered = err.lower() + # Ignore missing volumes. + if "no such volume" not in lowered: + sys.stderr.write(err) + sys.stderr.write("Failed to remove volumes.\n") + return int(vrm.returncode) + + sys.stdout.write("lightmem_server stopped docker services.\n") + if purge_volumes: + sys.stdout.write("Removed volumes: lightmem-redis-data, lightmem-etcd-data\n") + return 0 + + +def _port_open(host: str, port: int) -> bool: + import socket + + try: + with socket.create_connection((host, port), timeout=0.2): + return True + except OSError: + return False + + +def _redis_ping(host: str, port: int) -> bool: + import socket + + payload = "*1\r\n$4\r\nPING\r\n".encode("utf-8") + try: + with socket.create_connection((host, port), timeout=0.5) as s: + s.settimeout(0.5) + s.sendall(payload) + data = s.recv(64) + return data.startswith(b"+PONG") + except OSError: + return False + + +def _etcd_health(host: str, port: int) -> bool: + import json + import urllib.request + + url = f"http://{host}:{port}/health" + try: + with urllib.request.urlopen(url, timeout=0.8) as resp: + raw = resp.read() + obj = json.loads(raw.decode("utf-8")) if raw else {} + # etcd returns {"health":"true"} + return str(obj.get("health", "")).lower() == "true" + except Exception: + return False + + +def _detect_advertise_host() -> str | None: + """Best-effort detect a non-loopback IPv4 address for cross-host access.""" + import socket + + # Common technique: create a UDP socket to a public IP; no packets need to be sent. + try: + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + try: + s.connect(("8.8.8.8", 80)) + ip = s.getsockname()[0] + finally: + s.close() + if ip and not ip.startswith("127."): + return ip + except Exception: + pass + + # Fallback: resolve hostname. + try: + ip = socket.gethostbyname(socket.gethostname()) + if ip and not ip.startswith("127."): + return ip + except Exception: + pass + + return None + + +def _wait_port(host: str, port: int, *, timeout_s: float = 15.0) -> bool: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + if _port_open(host, port): + return True + time.sleep(0.1) + return False + + +def _start_redis_local(*, host: str, port: int, data_dir: str) -> subprocess.Popen: + redis_server = _which("redis-server") + if not redis_server: + raise RuntimeError("redis-server not found in PATH") + + if _port_open(host, port): + raise RuntimeError(f"redis port already in use: {host}:{port}") + + cmd = [ + redis_server, + "--bind", + host, + "--port", + str(port), + "--save", + "", + "--appendonly", + "yes", + "--appendfsync", + "everysec", + "--dir", + data_dir, + ] + proc = subprocess.Popen(cmd) + if not _wait_port(host, port, timeout_s=10.0): + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + raise RuntimeError("redis-server failed to start") + return proc + + +def _start_etcd_local(*, host: str, client_port: int, peer_port: int, data_dir: str) -> subprocess.Popen: + etcd = _which("etcd") + if not etcd: + raise RuntimeError("etcd not found in PATH") + + if _port_open(host, client_port): + raise RuntimeError(f"etcd client port already in use: {host}:{client_port}") + if _port_open(host, peer_port): + raise RuntimeError(f"etcd peer port already in use: {host}:{peer_port}") + + name = "coord" + cmd = [ + etcd, + "--name", + name, + "--data-dir", + data_dir, + "--listen-client-urls", + f"http://{host}:{client_port}", + "--advertise-client-urls", + f"http://{host}:{client_port}", + "--listen-peer-urls", + f"http://{host}:{peer_port}", + "--initial-advertise-peer-urls", + f"http://{host}:{peer_port}", + "--initial-cluster", + f"{name}=http://{host}:{peer_port}", + "--initial-cluster-state", + "new", + ] + proc = subprocess.Popen(cmd) + if not _wait_port(host, client_port, timeout_s=15.0): + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + raise RuntimeError("etcd failed to start") + return proc + + +def _run_local_services(*, index_port: int, coord_port: int, coord_peer_port: int) -> int: + host = "127.0.0.1" + with tempfile.TemporaryDirectory(prefix="lightmem_server_local_") as td: + redis_dir = os.path.join(td, "redis") + etcd_dir = os.path.join(td, "etcd") + os.makedirs(redis_dir, exist_ok=True) + os.makedirs(etcd_dir, exist_ok=True) + + redis_proc: subprocess.Popen | None = None + etcd_proc: subprocess.Popen | None = None + manage_redis = True + manage_etcd = True + + def _stop_all() -> None: + nonlocal redis_proc, etcd_proc + for p in (etcd_proc, redis_proc): + if p is None: + continue + try: + p.terminate() + except Exception: + pass + for p in (etcd_proc, redis_proc): + if p is None: + continue + try: + p.wait(timeout=8) + except Exception: + try: + p.kill() + except Exception: + pass + + def _sig_handler(_signum, _frame) -> None: # pragma: no cover + _stop_all() + raise SystemExit(0) + + signal.signal(signal.SIGINT, _sig_handler) + signal.signal(signal.SIGTERM, _sig_handler) + + try: + # If ports are already occupied, assume external services and reuse them. + if _port_open(host, index_port): + manage_redis = False + if not _redis_ping(host, index_port): + raise RuntimeError(f"port {host}:{index_port} is in use but does not look like Redis") + else: + redis_proc = _start_redis_local(host=host, port=index_port, data_dir=redis_dir) + + if _port_open(host, coord_port): + manage_etcd = False + if not _etcd_health(host, coord_port): + raise RuntimeError(f"port {host}:{coord_port} is in use but does not look like etcd") + else: + etcd_proc = _start_etcd_local( + host=host, + client_port=coord_port, + peer_port=coord_peer_port, + data_dir=etcd_dir, + ) + except Exception as e: + if manage_etcd or manage_redis: + _stop_all() + sys.stderr.write(f"Failed to start local services: {e}\n") + sys.stderr.write("Hint: install redis-server and etcd, or use --mode docker with docker compose.\n") + return 2 + + sys.stdout.write("lightmem_server local services ready (foreground).\n") + if not manage_redis: + sys.stdout.write("Index: (reused existing)\n") + if not manage_etcd: + sys.stdout.write("Coord: (reused existing)\n") + sys.stdout.write(f"Index: {host}:{index_port}\n") + sys.stdout.write(f"Coord: {host}:{coord_port}\n") + sys.stdout.write("Press Ctrl-C to stop.\n\n") + + # Block until a managed child exits (if we started any). + try: + while True: + if manage_redis and redis_proc is not None and redis_proc.poll() is not None: + sys.stderr.write("redis-server exited; stopping...\n") + break + if manage_etcd and etcd_proc is not None and etcd_proc.poll() is not None: + sys.stderr.write("etcd exited; stopping...\n") + break + time.sleep(0.2) + finally: + # Only stop what we started. + if manage_etcd or manage_redis: + _stop_all() + return 0 + + +def _compose_yaml(*, redis_port: int, etcd_client_port: int, etcd_peer_port: int) -> str: + # Minimal single-node index backend + coordinator backend. + # This is intentionally a thin wrapper that starts dependency services. + return dedent( + f""" + services: + index: + image: redis:7-alpine + container_name: lightmem-index + ports: + - "{redis_port}:6379" + command: ["redis-server", "--save", "", "--appendonly", "yes", "--appendfsync", "everysec"] + volumes: + - redis-data:/data + restart: unless-stopped + + coord: + image: quay.io/coreos/etcd:v3.5.12 + container_name: lightmem-coord + environment: + - ETCD_NAME=coord + - ETCD_DATA_DIR=/etcd-data + - ETCD_LISTEN_CLIENT_URLS=http://0.0.0.0:2379 + - ETCD_ADVERTISE_CLIENT_URLS=http://coord:2379 + - ETCD_LISTEN_PEER_URLS=http://0.0.0.0:2380 + - ETCD_INITIAL_ADVERTISE_PEER_URLS=http://coord:2380 + - ETCD_INITIAL_CLUSTER=coord=http://coord:2380 + - ETCD_INITIAL_CLUSTER_STATE=new + - ETCD_INITIAL_CLUSTER_TOKEN=lightmem + ports: + - "{etcd_client_port}:2379" + - "{etcd_peer_port}:2380" + volumes: + - etcd-data:/etcd-data + restart: unless-stopped + + volumes: + redis-data: + name: lightmem-redis-data + etcd-data: + name: lightmem-etcd-data + """ + ).lstrip() + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + prog="lightmem_server", + description="Start LightMem dependency services (index backend + coordinator backend) via Docker Compose.", + ) + parser.add_argument("--index-port", type=int, default=6379, help="Host port for the index backend (mapped to 6379)") + parser.add_argument("--coord-port", type=int, default=2379, help="Host port for the coordinator client (mapped to 2379)") + parser.add_argument("--coord-peer-port", type=int, default=2380, help="Host port for the coordinator peer (mapped to 2380)") + parser.add_argument( + "--mode", + choices=["auto", "docker", "local"], + default="auto", + help="Start mode: docker (docker compose), local (redis-server+etcd), or auto.", + ) + parser.add_argument( + "--stop", + action="store_true", + help="Stop docker-mode services (removes containers lightmem-index/lightmem-coord).", + ) + parser.add_argument( + "--purge-volumes", + action="store_true", + help="With --stop: also remove named volumes (lightmem-redis-data/lightmem-etcd-data).", + ) + parser.add_argument( + "--coord-ttl", + type=int, + default=None, + help="Coordinator TTL seconds (passed into PyLocalCacheService example)", + ) + + parser.add_argument( + "--coord-reconcile-sec", + type=float, + default=None, + help="Coordinator fallback reconcile interval seconds (passed into PyLocalCacheService example)", + ) + + parser.add_argument( + "--advertise-host", + default="", + help="Host/IP to print in endpoints for cross-host clients. Empty means auto-detect (docker) or 127.0.0.1 (local).", + ) + + args = parser.parse_args(argv) + + if args.purge_volumes and not args.stop: + sys.stderr.write("--purge-volumes requires --stop.\n") + return 2 + + if args.stop: + return _stop_docker_services(purge_volumes=bool(args.purge_volumes)) + + # The services started by this CLI may be reachable from other machines (docker mode). + # This only affects printed endpoints/snippets; it does not change bind behavior. + if args.mode == "local": + advertise_host = "127.0.0.1" + else: + advertise_host = str(args.advertise_host).strip() or (_detect_advertise_host() or "127.0.0.1") + + if args.mode == "local": + return _run_local_services(index_port=int(args.index_port), coord_port=int(args.coord_port), coord_peer_port=int(args.coord_peer_port)) + + compose = _docker_compose_cmd() + if args.mode in ("auto", "docker") and compose: + pass + elif args.mode == "docker": + sys.stderr.write("docker compose is required for lightmem_server --mode docker.\n") + sys.stderr.write("On Ubuntu: install docker-compose-plugin or docker-compose.\n") + return 2 + else: + # auto fallback + return _run_local_services(index_port=int(args.index_port), coord_port=int(args.coord_port), coord_peer_port=int(args.coord_peer_port)) + + if not (1 <= args.index_port <= 65535): + sys.stderr.write("Invalid --index-port.\n") + return 2 + if not (1 <= args.coord_port <= 65535): + sys.stderr.write("Invalid --coord-port.\n") + return 2 + if not (1 <= args.coord_peer_port <= 65535): + sys.stderr.write("Invalid --coord-peer-port.\n") + return 2 + + if args.coord_port == args.coord_peer_port: + sys.stderr.write("--coord-port and --coord-peer-port must differ.\n") + return 2 + + if args.coord_ttl is not None and args.coord_ttl <= 0: + sys.stderr.write("Invalid --coord-ttl (must be > 0).\n") + return 2 + if args.coord_reconcile_sec is not None and args.coord_reconcile_sec <= 0: + sys.stderr.write("Invalid --coord-reconcile-sec (must be > 0).\n") + return 2 + + + yml = _compose_yaml(redis_port=args.index_port, etcd_client_port=args.coord_port, etcd_peer_port=args.coord_peer_port) + + with tempfile.TemporaryDirectory(prefix="lightmem_server_") as td: + path = os.path.join(td, "docker-compose.lightmem-server.yml") + with open(path, "w", encoding="utf-8") as f: + f.write(yml) + + # Use a stable project name: we use stable container/volume names and a + # temp compose file path, so the default derived project name would vary + # per run and cause warnings about existing volumes. + cmd = [*compose, "--project-name", "lightmem_server", "-f", path, "up", "-d"] + rc = _run(cmd) + if rc != 0: + sys.stderr.write("Failed to start services via Docker Compose.\n") + return rc + + # Print minimal connection info for users. + sys.stdout.write("lightmem_server started dependency services.\n") + sys.stdout.write(f"Index: {advertise_host}:{args.index_port}\n") + sys.stdout.write(f"Coord: {advertise_host}:{args.coord_port}\n") + sys.stdout.write("\nExample Python for LightMem clients:\n") + sys.stdout.write(" from light_mem import PyLocalCacheService\n") + sys.stdout.write(" # ... prepare kvcache_tensor, file, etc ...\n") + sys.stdout.write(" svc = PyLocalCacheService(\n") + sys.stdout.write(" kvcache_tensor=kvcache_tensor,\n") + sys.stdout.write(" file=file,\n") + if args.index_port == 6379 and args.coord_port == 2379: + sys.stdout.write(f" index_endpoint=\"{advertise_host}\",\n") + else: + sys.stdout.write(f" index_endpoint=\"{advertise_host}:{args.index_port}\",\n") + sys.stdout.write(f" coord_endpoints=\"{advertise_host}:{args.coord_port}\",\n") + if args.coord_ttl is not None: + sys.stdout.write(f" coord_ttl={int(args.coord_ttl)},\n") + if args.coord_reconcile_sec is not None: + sys.stdout.write(f" coord_reconcile_sec={float(args.coord_reconcile_sec)},\n") + sys.stdout.write(" )\n") + sys.stdout.write("\n") + return 0 diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt old mode 100644 new mode 100755 index 64018fd..056cbea --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -16,19 +16,26 @@ list(APPEND CMAKE_PREFIX_PATH "${_torch_cmake_path}") message(STATUS "CMAKE_PREFIX_PATH: ${CMAKE_PREFIX_PATH}") find_package(Boost REQUIRED) +find_package(ZLIB REQUIRED) find_package(Torch REQUIRED) find_library(TORCH_PYTHON_LIBRARY torch_python PATH "${TORCH_INSTALL_PREFIX}/lib") # sources list set(source - pybind.cpp) + pybind.cpp + storage/local_cache_index.cpp + storage/local_storage_engine_public.cpp + storage/local_storage_engine_shard.cpp + storage/local_storage_engine_files.cpp + storage/local_storage_engine_journal.cpp + storage/local_storage_engine_recovery.cpp) # library Python_add_library(${name} MODULE ${source} WITH_SOABI) target_include_directories(${name} PRIVATE .) target_compile_definitions(${name} PRIVATE MODULE_NAME=${name}) target_compile_features(${name} PRIVATE cxx_std_17) -target_link_libraries(${name} PRIVATE Python::Python Boost::boost ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY}) +target_link_libraries(${name} PRIVATE Python::Python Boost::boost ZLIB::ZLIB ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY}) # Fix for macOS ARM64: use flat namespace to resolve symbols properly if(APPLE AND CMAKE_SYSTEM_PROCESSOR MATCHES "arm64") diff --git a/src/config.h b/src/config.h old mode 100644 new mode 100755 diff --git a/src/core/cache_task.h b/src/core/cache_task.h old mode 100644 new mode 100755 index 3415668..3d24bd1 --- a/src/core/cache_task.h +++ b/src/core/cache_task.h @@ -4,6 +4,7 @@ #include #include +#include #include #include #include @@ -69,6 +70,12 @@ class CacheTask : public std::enable_shared_from_this { throw std::invalid_argument("Invalid mode string. Use 'r' for Read or 'w' for Write."); } + // Timestamp (steady clock ticks) for end-to-end task latency measurement. + // This matches Python-side timing that measures create() -> ready(). + const auto now_duration = std::chrono::steady_clock::now().time_since_epoch(); + const int64_t now_ticks = static_cast(now_duration.count()); + submit_time_ticks.store(now_ticks, std::memory_order_relaxed); + blocks.reserve(hashs.size()); } @@ -112,6 +119,11 @@ class CacheTask : public std::enable_shared_from_this { Mode operation_mode; std::atomic completion_notified; std::vector page_already_list; + + // End-to-end timing hooks. + // NOTE: These are best-effort for logging/metrics, not correctness. + std::atomic submit_time_ticks{0}; + std::atomic finish_time_ticks{0}; }; } // namespace cache::task diff --git a/src/core/error.h b/src/core/error.h old mode 100644 new mode 100755 diff --git a/src/core/task_queue.h b/src/core/task_queue.h old mode 100644 new mode 100755 diff --git a/src/pybind.cpp b/src/pybind.cpp old mode 100644 new mode 100755 index 7b2cd61..9f52661 --- a/src/pybind.cpp +++ b/src/pybind.cpp @@ -31,11 +31,21 @@ PYBIND11_MODULE(MODULE_NAME, m) { .def("get_page_already_list", &CacheTask::get_page_already_list, "Get list of page indices already on disk"); py::class_(m, "LocalCacheService") - .def(py::init(), + .def(py::init(), py::arg("file"), py::arg("storage_size"), py::arg("num_of_shard"), py::arg("kvcache"), - py::arg("num_workers")) + py::arg("num_workers"), py::arg("index_endpoint") = std::string(), py::arg("bandwidth_log") = true, + py::arg("index_prefix") = std::string()) .def("run", &LocalCacheService::run) .def("query", &LocalCacheService::query) + .def("update_shard_assignments", &LocalCacheService::update_shard_assignments, py::arg("shard_ids"), + py::arg("epochs"), py::arg("draining")) + .def("recover_shard_to_redis", &LocalCacheService::recover_shard_to_redis, py::arg("shard_id")) + .def("recover_shard_to_redis_smart", &LocalCacheService::recover_shard_to_redis_smart, py::arg("shard_id")) + .def("shard_inflight", &LocalCacheService::shard_inflight, py::arg("shard_id")) + .def("shard_eviction_count", &LocalCacheService::shard_eviction_count, py::arg("shard_id")) + .def("eviction_count", &LocalCacheService::eviction_count) + .def("eviction_observed", &LocalCacheService::eviction_observed) .def("abort", &LocalCacheService::abort_task) .def("create", &LocalCacheService::create) .def("active_create_count", &LocalCacheService::active_create_count, py::arg("mode")) diff --git a/src/service/cache_service.h b/src/service/cache_service.h old mode 100644 new mode 100755 index da13d3f..eb44c8f --- a/src/service/cache_service.h +++ b/src/service/cache_service.h @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -60,8 +61,7 @@ struct alignas(8) CacheParam_t { CacheParam_t() : base_ptr(nullptr), page_size(0), page_stride(0), num_of_page(0) {} CacheParam_t(char *ptr, int64_t page_size, int64_t page_strd, int64_t pages) - : base_ptr(ptr), page_size(page_size), page_stride(page_strd), - num_of_page(pages) {} + : base_ptr(ptr), page_size(page_size), page_stride(page_strd), num_of_page(pages) {} }; // Generic Cache Service Base Class @@ -109,8 +109,7 @@ class CacheService { page_stride = strides[0] * element_size; const int64_t inferred_page_size = sizes[1] * element_size; - cache_info_ = - CacheParam_t((char *)kvcache.data_ptr(), inferred_page_size, page_stride, num_pages); + cache_info_ = CacheParam_t((char *)kvcache.data_ptr(), inferred_page_size, page_stride, num_pages); const int64_t page_bytes = cache_info_.page_size; const int64_t block_limit = resolve_max_block_size_bytes(); @@ -122,12 +121,13 @@ class CacheService { } block_size_ = page_bytes * pages_per_block; - printf("CacheService created with following cache info: \n"); - printf("\tNum of page: %lld \n", static_cast(cache_info_.num_of_page)); - printf("\tPage Size: %lld \n", static_cast(cache_info_.page_size)); - printf("\tPage Stride: %lld \n", static_cast(cache_info_.page_stride)); - printf("\tPages Per Block: %lld \n", static_cast(pages_per_block)); - printf("\tBlock Size: %lld \n", static_cast(block_size_)); + std::fprintf(stderr, "[light_mem] CacheService created with following cache info:\n"); + std::fprintf(stderr, "\tNum of page: %lld\n", static_cast(cache_info_.num_of_page)); + std::fprintf(stderr, "\tPage Size: %lld\n", static_cast(cache_info_.page_size)); + std::fprintf(stderr, "\tPage Stride: %lld\n", static_cast(cache_info_.page_stride)); + std::fprintf(stderr, "\tPages Per Block: %lld\n", static_cast(pages_per_block)); + std::fprintf(stderr, "\tBlock Size: %lld\n", static_cast(block_size_)); + std::fflush(stderr); } /** @@ -139,6 +139,11 @@ class CacheService { */ virtual std::vector query(const std::vector &hashs) = 0; + // Whether this service instance is running in online/distributed mode. + // In online mode shard ownership can be dynamic and certain optimizations (e.g. pre-query on write) + // must be disabled. + virtual bool online_mode() const { return false; } + int64_t block_size() const; int64_t page_size() const; @@ -203,17 +208,24 @@ class CacheService { // For write mode, query which pages are already in disk cache if (mode == "w") { - std::vector query_result = query(hashs); - const int32_t *page_ptr = reinterpret_cast(kv_page_indexer.data_ptr()); - - for (int64_t block_idx = 0; block_idx < static_cast(hashs.size()); ++block_idx) { - if (query_result[block_idx]) { - // This block is already in disk cache, add its page indices to page_already_list - const int64_t start_page_in_block = block_idx * page_per_block; - const int64_t end_page_in_block = std::min(start_page_in_block + page_per_block, num_of_pages); - - for (int64_t page_idx = start_page_in_block; page_idx < end_page_in_block; ++page_idx) { - task->page_already_list.push_back(page_ptr[page_idx]); + // NOTE: + // - In single-node mode this pre-query is cheap (O(1) shard lookup) and saves writes. + // - In online/distributed mode, shard ownership can be dynamic and pre-query tends to be expensive + // (e.g., scanning all shards and/or Redis RTT per hash), which can dominate write throughput. + // Dedupe/consistency is already enforced by the write path (WAL/Redis), so we skip this step. + if (!online_mode()) { + std::vector query_result = query(hashs); + const int32_t *page_ptr = reinterpret_cast(kv_page_indexer.data_ptr()); + + for (int64_t block_idx = 0; block_idx < static_cast(hashs.size()); ++block_idx) { + if (query_result[block_idx]) { + // This block is already in disk cache, add its page indices to page_already_list + const int64_t start_page_in_block = block_idx * page_per_block; + const int64_t end_page_in_block = std::min(start_page_in_block + page_per_block, num_of_pages); + + for (int64_t page_idx = start_page_in_block; page_idx < end_page_in_block; ++page_idx) { + task->page_already_list.push_back(page_ptr[page_idx]); + } } } } @@ -350,6 +362,13 @@ inline void CacheService::finalize_task(const std::shared_ptr(now_duration.count()); + task->finish_time_ticks.store(now_ticks, std::memory_order_relaxed); + } + std::atomic *active_counter = (task->operation_mode == cache::task::Mode::Read) ? &active_read_creates_ : &active_write_creates_; active_counter->fetch_sub(1, std::memory_order_relaxed); diff --git a/src/service/local_cache_service.h b/src/service/local_cache_service.h old mode 100644 new mode 100755 index dd3196c..c5066eb --- a/src/service/local_cache_service.h +++ b/src/service/local_cache_service.h @@ -11,7 +11,10 @@ #include #include #include +#include +#include #include +#include #include #include #include @@ -45,11 +48,10 @@ class LocalCacheService : public CacheService { * @param num_workers Number of worker threads */ LocalCacheService(const string &file, size_t storage_size, size_t num_shard, const torch::Tensor &kvcache, - const size_t num_workers) - : CacheService(kvcache), stop_(false), num_workers_(num_workers), block_size_(0), total_written_bytes_(0), - first_write_time_ticks_(0), last_write_time_ticks_(0), last_log_time_(), last_logged_bytes_(0), - total_read_bytes_(0), first_read_time_ticks_(0), last_read_time_ticks_(0), read_last_log_time_(), - read_last_logged_bytes_(0) { + const size_t num_workers, const std::string &index_endpoint = "", bool bandwidth_log = true, + const std::string &index_prefix = "") + : CacheService(kvcache), stop_(false), num_workers_(num_workers), block_size_(0), bandwidth_log_(bandwidth_log), + online_mode_(!index_endpoint.empty()) { block_size_ = static_cast(this->block_size()); if (storage_size < block_size_) { @@ -58,7 +60,7 @@ class LocalCacheService : public CacheService { ensure_disk_capacity(file, storage_size, num_shard); - storage_ = make_unique(file, storage_size, num_shard, block_size_); + storage_ = make_unique(file, storage_size, num_shard, block_size_, index_endpoint, index_prefix); // Use unique_ptr for exception safety - if any allocation fails, previous allocations are automatically cleaned up r_cpu_buffers_.reserve(num_workers_); @@ -90,11 +92,97 @@ class LocalCacheService : public CacheService { * This function will throw no exception or error. */ std::vector query(const std::vector &hashs) override { - std::vector ret; - ret.reserve(hashs.size()); - std::transform(hashs.begin(), hashs.end(), std::back_inserter(ret), - [this](const auto &hash) { return storage_->query(hash); }); - return ret; + if (!storage_) { + return std::vector(hashs.size(), false); + } + return storage_->queryMany(hashs); + } + + // Online/distributed mode hooks (optionally driven by a Python coordinator). + // These APIs are control-plane helpers for shard ownership handoff. + // Notes: + // - "writable" controls whether new writes are allowed for a shard. + // - "draining" is a transition state: shard is still considered owned/readable, but new writes are rejected. + // - "epoch" is a generation number used to fence writes across ownership changes. + + /** + * @brief Update this node's shard assignment state in online/distributed mode. + * + * @param shard_ids Shard IDs owned/managed by this node. + * @param epochs Per-shard epoch (generation). Must align 1:1 with shard_ids. + * @param draining Per-shard draining flag (0/1). If 1, shard is treated as not writable (no new writes). + * + * Requirements: shard_ids/epochs/draining must have the same length. + */ + void update_shard_assignments(const std::vector &shard_ids, const std::vector &epochs, + const std::vector &draining) { + if (!storage_) { + return; + } + storage_->updateShardAssignments(shard_ids, epochs, draining); + } + + /** + * @brief Rebuild Redis index for a shard from local snapshot + WAL. + * + * Intended to be called when a node (re-)acquires write permission for a shard. + */ + void recover_shard_to_redis(size_t shard_id) { + if (!storage_) { + return; + } + storage_->recoverShardToRedis(shard_id); + } + + /** + * @brief Recover Redis index for a shard using smart policy (full vs incremental). + */ + void recover_shard_to_redis_smart(size_t shard_id) { + if (!storage_) { + return; + } + storage_->recoverShardToRedisSmart(shard_id); + } + + /** + * @brief Observability: current in-flight write operations for a shard. + * + * This is a best-effort counter used by the coordinator/control-plane to decide when a shard is drained. + */ + uint32_t shard_inflight(size_t shard_id) const { + if (!storage_) { + return 0; + } + return storage_->shardInflight(shard_id); + } + + uint64_t shard_written_bytes(size_t shard_id) const { + if (!storage_) { + return 0; + } + return storage_->shardWrittenBytes(shard_id); + } + + // True LRU eviction observability (local to this node). + uint64_t shard_eviction_count(size_t shard_id) const { + if (!storage_) { + return 0; + } + return storage_->shardEvictionCount(shard_id); + } + + uint64_t eviction_count() const { + if (!storage_) { + return 0; + } + return storage_->evictionCount(); + } + + bool eviction_observed() const { + if (!storage_) { + return false; + } + return storage_->evictionObserved(); } /** @@ -108,8 +196,29 @@ class LocalCacheService : public CacheService { } protected: + bool online_mode() const override { return online_mode_; } + void on_task_finalized(const std::shared_ptr &task) override { + if (!bandwidth_log_) { + return; + } + if (task->operation_mode == cache::task::Mode::Write) { + // Keep log format identical to historical output, but compute recent write speed + // with end-to-end task timing to match benchmark scripts more closely. + // - bytes (for speed) = pages * page_size (logical/requested bytes) + // - time (for speed) = create() -> ready() end-to-end (per task) + const int64_t start_ticks = task->submit_time_ticks.load(std::memory_order_relaxed); + const int64_t end_ticks = task->finish_time_ticks.load(std::memory_order_relaxed); + if (start_ticks != 0 && end_ticks != 0 && end_ticks > start_ticks) { + const int64_t pages = task->page_indexer.numel(); + const uint64_t logical_bytes = + static_cast(pages) * static_cast(this->page_size()); + const int64_t elapsed_ticks = end_ticks - start_ticks; + window_logical_written_bytes_.fetch_add(logical_bytes, std::memory_order_relaxed); + window_logical_write_time_ticks_.fetch_add(elapsed_ticks, std::memory_order_relaxed); + } + // Try to acquire the lock, skip logging if contention occurs std::unique_lock guard(log_mutex_, std::try_to_lock); if (!guard.owns_lock()) { @@ -133,22 +242,23 @@ class LocalCacheService : public CacheService { const uint64_t delta_bytes = (total >= previous_bytes) ? (total - previous_bytes) : 0; double speed_gbps = 0.0; - // Use actual I/O time for speed calculation - if (last_log_time_ != std::chrono::steady_clock::time_point{}) { - const double elapsed_sec = - std::chrono::duration_cast>(now - last_log_time_).count(); - if (elapsed_sec > 0.0) { - speed_gbps = (static_cast(delta_bytes) / (1024.0 * 1024.0 * 1024.0)) / elapsed_sec; + // Prefer end-to-end (script-like) accounting for speed. + const uint64_t window_bytes = window_logical_written_bytes_.exchange(0, std::memory_order_relaxed); + const int64_t window_ticks = window_logical_write_time_ticks_.exchange(0, std::memory_order_relaxed); + if (window_bytes != 0 && window_ticks > 0) { + const auto window_dur = + std::chrono::steady_clock::duration(static_cast(window_ticks)); + const double window_sec = std::chrono::duration_cast>(window_dur).count(); + if (window_sec > 0.0) { + speed_gbps = (static_cast(window_bytes) / (1024.0 * 1024.0 * 1024.0)) / window_sec; } - } else { - // For the first log, use time from first write start to now - const int64_t first_ticks = first_write_time_ticks_.load(std::memory_order_relaxed); - if (first_ticks != 0) { - const auto first_duration = - std::chrono::steady_clock::duration(static_cast(first_ticks)); - const auto first_time = std::chrono::steady_clock::time_point(first_duration); + } + + // Fallback: if window accounting isn't available, keep legacy behavior. + if (speed_gbps == 0.0) { + if (last_log_time_ != std::chrono::steady_clock::time_point{}) { const double elapsed_sec = - std::chrono::duration_cast>(now - first_time).count(); + std::chrono::duration_cast>(now - last_log_time_).count(); if (elapsed_sec > 0.0) { speed_gbps = (static_cast(delta_bytes) / (1024.0 * 1024.0 * 1024.0)) / elapsed_sec; } @@ -159,63 +269,57 @@ class LocalCacheService : public CacheService { last_logged_bytes_ = total; const double total_gb = static_cast(total) / (1024.0 * 1024.0 * 1024.0); - printf("[light_mem] cumulative disk write size: %.2f GB, recent write speed: %.2f GB/s\n", total_gb, speed_gbps); + std::fprintf(stderr, "[light_mem] cumulative disk write size: %.2f GB, recent write speed: %.2f GB/s\n", + total_gb, speed_gbps); + std::fflush(stderr); return; } if (task->operation_mode != cache::task::Mode::Read) { return; } + + // For read mode, align bandwidth accounting with benchmark scripts: + // - bytes (for speed) = pages * page_size (logical/requested bytes) + // - time (for speed) = create() -> ready() end-to-end (per task) + const int64_t start_ticks = task->submit_time_ticks.load(std::memory_order_relaxed); + const int64_t end_ticks = task->finish_time_ticks.load(std::memory_order_relaxed); + if (start_ticks != 0 && end_ticks != 0 && end_ticks > start_ticks) { + const int64_t pages = task->page_indexer.numel(); + const uint64_t logical_bytes = + static_cast(pages) * static_cast(this->page_size()); + const int64_t elapsed_ticks = end_ticks - start_ticks; + window_logical_read_bytes_.fetch_add(logical_bytes, std::memory_order_relaxed); + window_logical_read_time_ticks_.fetch_add(elapsed_ticks, std::memory_order_relaxed); + } + if (active_read_creates_.load(std::memory_order_relaxed) != 0) { return; } std::lock_guard guard(read_log_mutex_); - const uint64_t total_read = total_read_bytes_.load(std::memory_order_relaxed); - if (total_read == 0) { + const uint64_t window_bytes = window_logical_read_bytes_.exchange(0, std::memory_order_relaxed); + const int64_t window_ticks = window_logical_read_time_ticks_.exchange(0, std::memory_order_relaxed); + if (window_bytes == 0 || window_ticks <= 0) { return; } - // Calculate batch read amount (since last queue empty) - const uint64_t previous_read = read_last_logged_bytes_; - const uint64_t delta_read = (total_read >= previous_read) ? (total_read - previous_read) : 0; - if (delta_read == 0) { + const auto window_dur = + std::chrono::steady_clock::duration(static_cast(window_ticks)); + const double window_sec = std::chrono::duration_cast>(window_dur).count(); + if (window_sec <= 0.0) { return; } - // Use actual I/O time: from first read start to last read completion - const int64_t first_ticks = first_read_time_ticks_.load(std::memory_order_relaxed); - const int64_t last_ticks = last_read_time_ticks_.load(std::memory_order_relaxed); - - double speed_gbps = 0.0; - if (first_ticks != 0 && last_ticks != 0 && last_ticks > first_ticks) { - const auto first_duration = - std::chrono::steady_clock::duration(static_cast(first_ticks)); - const auto last_duration = - std::chrono::steady_clock::duration(static_cast(last_ticks)); - const auto first_time = std::chrono::steady_clock::time_point(first_duration); - const auto last_time = std::chrono::steady_clock::time_point(last_duration); - const double elapsed_sec = - std::chrono::duration_cast>(last_time - first_time).count(); - if (elapsed_sec > 0.0) { - speed_gbps = (static_cast(delta_read) / (1024.0 * 1024.0 * 1024.0)) / elapsed_sec; - } - } - + const double window_gb = static_cast(window_bytes) / (1024.0 * 1024.0 * 1024.0); + const double speed_gbps = window_gb / window_sec; if (speed_gbps <= 0.0) { return; } - const double delta_read_gb = static_cast(delta_read) / (1024.0 * 1024.0 * 1024.0); - - // Reset counters for next batch after queue is empty - read_last_log_time_ = std::chrono::steady_clock::now(); - read_last_logged_bytes_ = 0; // Reset to 0 instead of total_read - total_read_bytes_.store(0, std::memory_order_relaxed); // Clear accumulated bytes - first_read_time_ticks_.store(0, std::memory_order_relaxed); // Reset timing - last_read_time_ticks_.store(0, std::memory_order_relaxed); // Reset timing - printf("[light_mem] batch read size: %.2f GB, read speed: %.2f GB/s\n", delta_read_gb, speed_gbps); + std::fprintf(stderr, "[light_mem] batch read size: %.2f GB, read speed: %.2f GB/s\n", window_gb, speed_gbps); + std::fflush(stderr); } private: @@ -356,6 +460,7 @@ class LocalCacheService : public CacheService { // 2. Temporary failure (all slots busy, I/O error) // 3. Write was skipped // This is acceptable for cache operations - treat as success to avoid abort + if (written == block_size_) { total_written_bytes_.fetch_add(static_cast(written), std::memory_order_relaxed); } @@ -388,6 +493,33 @@ class LocalCacheService : public CacheService { const int64_t total_pages = info.num_of_page; const int64_t page_bytes = page_size; + // Fast path: if destination pages are contiguous in memory and indices form a contiguous range, + // we can copy the entire span in one memcpy. This is the common case for benchmarks using + // kv_page_indexer=torch.arange(...). + if (num_of_page > 0 && page_stride == page_bytes) { + const int32_t first = page_idx[0]; + if (first < 0 || first >= total_pages) { + throw std::runtime_error("kv page index out of range in cpu_scatter."); + } + bool contiguous = true; + for (int64_t local_page = 1; local_page < num_of_page; ++local_page) { + const int32_t expected = first + static_cast(local_page); + if (page_idx[local_page] != expected) { + contiguous = false; + break; + } + } + if (contiguous) { + const int32_t last = first + static_cast(num_of_page - 1); + if (last < 0 || last >= total_pages) { + throw std::runtime_error("kv page index out of range in cpu_scatter."); + } + char *dst_ptr = info.base_ptr + static_cast(first) * page_stride; + std::memcpy(dst_ptr, block, static_cast(num_of_page) * static_cast(page_bytes)); + return; + } + } + for (int64_t local_page = 0; local_page < num_of_page; ++local_page) { const int32_t dst_page = page_idx[local_page]; if (dst_page < 0 || dst_page >= total_pages) { @@ -421,6 +553,32 @@ class LocalCacheService : public CacheService { const int64_t total_pages = info.num_of_page; const int64_t page_bytes = page_size; + // Fast path: if source pages are contiguous in memory and indices form a contiguous range, + // we can gather the entire span in one memcpy. + if (num_of_page > 0 && page_stride == page_bytes) { + const int32_t first = page_idx[0]; + if (first < 0 || first >= total_pages) { + throw std::runtime_error("kv page index out of range in cpu_gather."); + } + bool contiguous = true; + for (int64_t local_page = 1; local_page < num_of_page; ++local_page) { + const int32_t expected = first + static_cast(local_page); + if (page_idx[local_page] != expected) { + contiguous = false; + break; + } + } + if (contiguous) { + const int32_t last = first + static_cast(num_of_page - 1); + if (last < 0 || last >= total_pages) { + throw std::runtime_error("kv page index out of range in cpu_gather."); + } + const char *src_ptr = info.base_ptr + static_cast(first) * page_stride; + std::memcpy(block, src_ptr, static_cast(num_of_page) * static_cast(page_bytes)); + return; + } + } + for (int64_t local_page = 0; local_page < num_of_page; ++local_page) { const int32_t src_page = page_idx[local_page]; if (src_page < 0 || src_page >= total_pages) { @@ -434,26 +592,33 @@ class LocalCacheService : public CacheService { } } - size_t block_size_; ///< Block size - unique_ptr storage_; ///< Local storage engine - vector workers_; ///< Worker threads - bool stop_; ///< Thread stop flag - size_t num_workers_; ///< Number of worker threads - vector> r_cpu_buffers_; ///< CPU buffers for read worker (RAII managed) - vector> w_cpu_buffers_; ///< CPU buffers for write worker (RAII managed) - std::atomic total_written_bytes_; ///< Total bytes written to disk - std::atomic first_write_time_ticks_; ///< First write start time in steady clock ticks - std::atomic last_write_time_ticks_; ///< Last write completion time in steady clock ticks - std::mutex log_mutex_; ///< Protects write rate reporting - std::chrono::steady_clock::time_point last_log_time_; ///< Last write log timestamp - uint64_t last_logged_bytes_; ///< Bytes recorded at last write log - - std::atomic total_read_bytes_; ///< Total bytes read from disk - std::atomic first_read_time_ticks_; ///< First read start time in steady clock ticks - std::atomic last_read_time_ticks_; ///< Last read completion time in steady clock ticks - std::mutex read_log_mutex_; ///< Protects read rate reporting - std::chrono::steady_clock::time_point read_last_log_time_; ///< Last read log timestamp - uint64_t read_last_logged_bytes_; ///< Bytes recorded at last read log + size_t block_size_; ///< Block size + unique_ptr storage_; ///< Local storage engine + vector workers_; ///< Worker threads + bool stop_; ///< Thread stop flag + size_t num_workers_; ///< Number of worker threads + vector> r_cpu_buffers_; ///< CPU buffers for read worker (RAII managed) + vector> w_cpu_buffers_; ///< CPU buffers for write worker (RAII managed) + + bool bandwidth_log_{true}; + bool online_mode_{false}; + + std::atomic total_written_bytes_{0}; + std::atomic window_logical_written_bytes_{0}; + std::atomic window_logical_write_time_ticks_{0}; + std::atomic first_write_time_ticks_{0}; + std::atomic last_write_time_ticks_{0}; + std::mutex log_mutex_; + std::chrono::steady_clock::time_point last_log_time_{}; + uint64_t last_logged_bytes_{0}; + + std::atomic total_read_bytes_{0}; + std::atomic window_logical_read_bytes_{0}; + std::atomic window_logical_read_time_ticks_{0}; + std::atomic first_read_time_ticks_{0}; + std::atomic last_read_time_ticks_{0}; + std::mutex read_log_mutex_; + uint64_t read_last_logged_bytes_{0}; // Ensure the backing storage path exposes enough disk capacity for the requested cache size. static void ensure_disk_capacity(const string &file, size_t storage_size, size_t num_shard) { diff --git a/src/storage/local_cache_index.cpp b/src/storage/local_cache_index.cpp new file mode 100755 index 0000000..106dc1d --- /dev/null +++ b/src/storage/local_cache_index.cpp @@ -0,0 +1,399 @@ +#include "storage/local_cache_index.h" + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace cache { +namespace storage { + +static const uint32_t SNAPSHOT_MAGIC = 0x534E4150; // SNAP +static const uint32_t SNAPSHOT_VERSION = 1; + +LocalCacheIndex::LocalCacheIndex(size_t capacity) : capacity_(capacity) { + // Allocate low offsets first. + // empty_block_list_.back() is used as the next free slot; so we push in reverse order + // and pop from the back to get 0, 1, 2, ... + for (size_t i = capacity; i > 0; --i) { + empty_block_list_.push_back(i - 1); + } +} + +void LocalCacheIndex::reset() { + std::lock_guard lock(index_lock_); + lru_list_.clear(); + index_.clear(); + empty_block_list_.clear(); + eviction_count_ = 0; + for (size_t i = capacity_; i > 0; --i) { + empty_block_list_.push_back(i - 1); + } +} + +void LocalCacheIndex::put_ready(const std::string &hash, size_t slot_id, uint32_t crc) { + std::lock_guard lock(index_lock_); + + if (slot_id >= capacity_) { + return; + } + + // If hash already exists, just update slot and promote. + auto it = index_.find(hash); + if (it != index_.end()) { + it->second.slot_id = slot_id; + it->second.ready = true; + it->second.writing = false; + it->second.crc = crc; + lru_list_.splice(lru_list_.begin(), lru_list_, it->second.lru_iterator); + return; + } + + // If slot is already used by someone else, evict that hash. + for (auto map_it = index_.begin(); map_it != index_.end(); ++map_it) { + if (map_it->second.slot_id == slot_id) { + lru_list_.erase(map_it->second.lru_iterator); + index_.erase(map_it); + break; + } + } + + // Ensure slot is not in empty list. + for (auto e = empty_block_list_.begin(); e != empty_block_list_.end(); ++e) { + if (*e == slot_id) { + empty_block_list_.erase(e); + break; + } + } + + // If over capacity (shouldn't happen if slot_id is from valid range), evict LRU. + if (index_.size() >= capacity_ && !lru_list_.empty()) { + const std::string victim = lru_list_.back(); + auto vit = index_.find(victim); + if (vit != index_.end()) { + size_t freed = vit->second.slot_id; + lru_list_.pop_back(); + index_.erase(vit); + empty_block_list_.push_back(freed); + } + } + + lru_list_.push_front(hash); + index_[hash] = {lru_list_.begin(), slot_id, true, false, crc}; +} + +bool LocalCacheIndex::exists(const std::string &hash) { + std::lock_guard lock(index_lock_); + auto it = index_.find(hash); + // Only return true if data is fully written and readable (not just allocated) + if (it == index_.end() || !it->second.ready || it->second.writing) { + return false; + } + lru_list_.splice(lru_list_.begin(), lru_list_, it->second.lru_iterator); + return true; +} + +int LocalCacheIndex::acquire_slot(const std::string &hash, size_t &slot_id, std::string &evicted_hash) { + std::lock_guard lock(index_lock_); + + auto existing = index_.find(hash); + if (existing != index_.end()) { + if (existing->second.writing) { + return -1; // Write in progress, caller should retry + } + lru_list_.splice(lru_list_.begin(), lru_list_, existing->second.lru_iterator); + slot_id = existing->second.slot_id; + evicted_hash.clear(); + return 0; // Already exists and ready + } + + if (!empty_block_list_.empty()) { + slot_id = empty_block_list_.back(); + empty_block_list_.pop_back(); + evicted_hash.clear(); + } else { + // Need to evict a victim from LRU list + // Find a victim that is not currently being written (writing=false) + auto it = lru_list_.end(); + bool found_victim = false; + + while (it != lru_list_.begin()) { + --it; + const std::string &candidate_hash = *it; + auto candidate_it = index_.find(candidate_hash); + + // Data structure inconsistency detected, return failure to avoid corruption + if (candidate_it == index_.end()) { + std::fprintf(stderr, "[light_mem error] LRU list and index inconsistent, returning failure\n"); + return -1; + } + + if (!candidate_it->second.writing) { + slot_id = candidate_it->second.slot_id; + evicted_hash = candidate_hash; + lru_list_.erase(it); + index_.erase(candidate_it); + eviction_count_++; + found_victim = true; + break; + } + } + + if (!found_victim) { + // All slots are busy writing or LRU list is empty + // This is temporary congestion, return -1 to let caller retry + return -1; + } + } + + lru_list_.push_front(hash); + // Mark as writing=true, ready=false + index_[hash] = {lru_list_.begin(), slot_id, false, true, 0}; + return 1; +} + +uint64_t LocalCacheIndex::eviction_count() const { + std::lock_guard lock(index_lock_); + return eviction_count_; +} + +bool LocalCacheIndex::eviction_observed() const { return eviction_count() > 0; } + +void LocalCacheIndex::mark_ready(const std::string &hash, uint32_t crc) { + std::lock_guard lock(index_lock_); + auto it = index_.find(hash); + if (it != index_.end()) { + it->second.ready = true; + it->second.writing = false; + if (crc != 0) { + it->second.crc = crc; + } + } +} + +void LocalCacheIndex::remove(const std::string &hash) { + std::lock_guard lock(index_lock_); + auto it = index_.find(hash); + if (it != index_.end()) { + size_t slot_id = it->second.slot_id; + lru_list_.erase(it->second.lru_iterator); + index_.erase(it); + empty_block_list_.push_back(slot_id); + } +} + +size_t LocalCacheIndex::get_offset(const std::string &hash) { + std::lock_guard lock(index_lock_); + auto it = index_.find(hash); + if (it == index_.end() || !it->second.ready || it->second.writing) { + return static_cast(-1); + } + return it->second.slot_id; +} + +bool LocalCacheIndex::get_offset_and_crc(const std::string &hash, size_t &slot_id, uint32_t &crc) { + std::lock_guard lock(index_lock_); + auto it = index_.find(hash); + if (it == index_.end() || !it->second.ready || it->second.writing) { + return false; + } + slot_id = it->second.slot_id; + crc = it->second.crc; + return true; +} + +void LocalCacheIndex::dump_ready(std::vector &out) { + std::lock_guard lock(index_lock_); + out.clear(); + out.reserve(index_.size()); + for (const auto &kv : index_) { + const auto &e = kv.second; + if (e.ready && !e.writing) { + out.emplace_back(kv.first); + } + } +} + +bool LocalCacheIndex::saveToSnapshot(const std::string &filename) { + std::lock_guard lock(index_lock_); + + std::string tmp_filename = filename + ".tmp"; + // Shared storage: allow other nodes/users to read+write snapshots. + int fd = ::open(tmp_filename.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0666); + if (fd < 0) { + return false; + } + + // Best-effort: umask may have reduced permission bits. + (void)::fchmod(fd, 0666); + + // Header: Magic(4) + Version(4) + Count(8) + uint32_t magic = SNAPSHOT_MAGIC; + uint32_t version = SNAPSHOT_VERSION; + uint64_t count = 0; + + // Count valid entries (ready=true) + for (const auto &kv : index_) { + if (kv.second.ready && !kv.second.writing) { + count++; + } + } + + auto write_all = [&](const void *p, size_t n) -> bool { + const char *buf = static_cast(p); + size_t left = n; + while (left > 0) { + ssize_t w = ::write(fd, buf, left); + if (w <= 0) { + return false; + } + buf += static_cast(w); + left -= static_cast(w); + } + return true; + }; + + if (!write_all(&magic, sizeof(magic)) || !write_all(&version, sizeof(version)) || !write_all(&count, sizeof(count))) { + ::close(fd); + std::remove(tmp_filename.c_str()); + return false; + } + + // Write entries in Cold-to-Hot order (reverse LRU) + // lru_list_: front=Hot, back=Cold + // We iterate from rbegin() (Cold) to rend() (Hot) + for (auto it = lru_list_.rbegin(); it != lru_list_.rend(); ++it) { + const std::string &hash = *it; + auto idx_it = index_.find(hash); + if (idx_it != index_.end()) { + const auto &entry = idx_it->second; + if (entry.ready && !entry.writing) { + uint32_t hash_len = static_cast(hash.size()); + uint64_t slot_id = static_cast(entry.slot_id); + uint32_t crc = static_cast(entry.crc); + + if (!write_all(&hash_len, sizeof(hash_len)) || !write_all(hash.data(), hash_len) || + !write_all(&slot_id, sizeof(slot_id)) || !write_all(&crc, sizeof(crc))) { + ::close(fd); + std::remove(tmp_filename.c_str()); + return false; + } + } + } + } + + // Ensure file contents are durable before rename. + if (::fsync(fd) != 0) { + ::close(fd); + std::remove(tmp_filename.c_str()); + return false; + } + ::close(fd); + + // Atomic replace. + if (std::rename(tmp_filename.c_str(), filename.c_str()) != 0) { + std::remove(tmp_filename.c_str()); + return false; + } + + // Best-effort: fsync directory to persist rename. + std::string dir = "."; + auto pos = filename.find_last_of('/'); + if (pos != std::string::npos) { + dir = filename.substr(0, pos); + } + int dfd = ::open(dir.c_str(), O_RDONLY | O_DIRECTORY); + if (dfd >= 0) { + (void)::fsync(dfd); + ::close(dfd); + } + + return true; +} + +bool LocalCacheIndex::loadFromSnapshot(const std::string &filename) { + // Note: Caller should ensure reset() is called before or after as needed, + // but typically we load into a fresh index. + // We use put_ready which is thread-safe and handles locking. + + int fd = ::open(filename.c_str(), O_RDONLY); + if (fd < 0) { + return false; + } + + auto read_all = [&](void *p, size_t n) -> bool { + char *buf = static_cast(p); + size_t left = n; + while (left > 0) { + ssize_t r = ::read(fd, buf, left); + if (r <= 0) { + return false; + } + buf += static_cast(r); + left -= static_cast(r); + } + return true; + }; + + uint32_t magic = 0; + uint32_t version = 0; + uint64_t count = 0; + + if (!read_all(&magic, sizeof(magic)) || !read_all(&version, sizeof(version)) || !read_all(&count, sizeof(count))) { + ::close(fd); + return false; + } + if (magic != SNAPSHOT_MAGIC || version != SNAPSHOT_VERSION) { + ::close(fd); + return false; + } + + for (uint64_t i = 0; i < count; i++) { + uint32_t hash_len = 0; + if (!read_all(&hash_len, sizeof(hash_len)) || hash_len > 4096) { + ::close(fd); + return false; + } + + std::string hash; + hash.resize(hash_len); + if (hash_len > 0 && !read_all(hash.data(), hash_len)) { + ::close(fd); + return false; + } + + uint64_t slot_id = 0; + if (!read_all(&slot_id, sizeof(slot_id))) { + ::close(fd); + return false; + } + + uint32_t crc = 0; + if (!read_all(&crc, sizeof(crc))) { + ::close(fd); + return false; + } + + // Insert into index (Hot-end insertion) + // Since we saved Cold-to-Hot, reading sequentially and inserting at Head + // will result in [Hot, ..., Cold] order in memory. + // Wait: + // Saved: Cold, ..., Hot + // Read 1 (Cold): put_ready -> [Cold] + // Read 2 (Medium): put_ready -> [Medium, Cold] + // Read 3 (Hot): put_ready -> [Hot, Medium, Cold] + // Correct. + put_ready(hash, static_cast(slot_id), crc); + } + + ::close(fd); + return true; +} + +} // namespace storage +} // namespace cache diff --git a/src/storage/local_cache_index.h b/src/storage/local_cache_index.h new file mode 100755 index 0000000..79f7423 --- /dev/null +++ b/src/storage/local_cache_index.h @@ -0,0 +1,95 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cache { +namespace storage { + +/** + * @brief Cache index class + */ +class LocalCacheIndex { +public: + /** + * @brief Internal structure to store LRU list iterator and exists call count + */ + struct IndexEntry { + std::list::iterator lru_iterator; + size_t slot_id; // Slot index + bool ready; // Data is written to disk and readable + bool writing; // Slot is allocated but disk write in progress (evictable but not readable) + uint32_t crc; // CRC of the data block (0 means unknown) + }; + + /** + * @brief Constructor + * @param capacity Maximum number of hash values to store + */ + explicit LocalCacheIndex(size_t capacity); + + void reset(); + + // Insert a ready mapping (used by recovery / redis warmup). + // If slot is already held by another hash, that entry is removed. + void put_ready(const std::string &hash, size_t slot_id, uint32_t crc = 0); + + /** + * @brief Checks if a hash value exists + * + * If it exists, returns true and updates the LRU order (moves to the front); + * If it doesn't exist, returns false. + */ + bool exists(const std::string &hash); + + /** + * @brief Acquire a slot for a hash, allocating or reusing storage as needed. + * + * @return 1 if newly inserted, 0 if already existed, -1 if temporarily failed (all slots busy) + */ + int acquire_slot(const std::string &hash, size_t &slot_id, std::string &evicted_hash); + + // True LRU eviction observability. + // This counts only evictions triggered by acquire_slot() due to full capacity. + uint64_t eviction_count() const; + bool eviction_observed() const; + + void mark_ready(const std::string &hash, uint32_t crc = 0); + void remove(const std::string &hash); + + /** + * @brief Gets the offset of a hash value + * @return slot id if exists and ready, size_t(-1) otherwise + */ + size_t get_offset(const std::string &hash); + + // Returns true iff hash exists and is ready; outputs slot and crc. + // crc==0 means CRC is unknown/unavailable. + bool get_offset_and_crc(const std::string &hash, size_t &slot_id, uint32_t &crc); + + // Snapshot operations + bool saveToSnapshot(const std::string &filename); + bool loadFromSnapshot(const std::string &filename); + + // Best-effort dump of ready entries for building auxiliary indices. + // Thread-safe snapshot under the internal mutex. + void dump_ready(std::vector &out); + +private: + size_t capacity_; ///< Maximum number of hash values to store + std::list lru_list_; ///< LRU list, head is most recent, tail is least recently used + std::list empty_block_list_; ///< List of free disk blocks + std::unordered_map index_; ///< Map from hash value to IndexEntry + mutable std::mutex index_lock_; ///< Mutex protecting index data structures + + uint64_t eviction_count_{0}; +}; + +} // namespace storage +} // namespace cache diff --git a/src/storage/local_storage_engine.h b/src/storage/local_storage_engine.h old mode 100644 new mode 100755 index 59f5722..83d4c8e --- a/src/storage/local_storage_engine.h +++ b/src/storage/local_storage_engine.h @@ -1,382 +1,218 @@ -#pragma once - -#include "core/error.h" -#include "storage/storage_engine.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace cache { -namespace storage { - -/** - * @brief Cache index class - */ -class LocalCacheIndex { -public: - /** - * @brief Constructor - * @param capacity Maximum number of hash values to store - */ - LocalCacheIndex(size_t capacity) : capacity_(capacity) { - for (size_t i = 0; i < capacity; i++) { - empty_block_list_.push_back(i); - } - } - - /** - * @brief Checks if a hash value exists - * - * If it exists, returns true and updates the LRU order (moves to the front); - * If it doesn't exist, returns false. - * - * @param hash The hash value to search for - * @return True if exists, false otherwise - */ - bool exists(const std::string &hash) { - std::lock_guard lock(index_lock_); - auto it = index_.find(hash); - // Only return true if data is fully written and readable (not just allocated) - if (it == index_.end() || !it->second.ready || it->second.writing) { - return false; - } - lru_list_.splice(lru_list_.begin(), lru_list_, it->second.lru_iterator); - return true; - } - - /** - * @brief Acquire a slot for a hash, allocating or reusing storage as needed. - * - * If the hash already exists, its offset is returned and the entry promoted to MRU. - * If the hash is new and free space remains, a free offset is consumed. - * Otherwise the LRU entry is evicted and its offset reused for the new hash. - * - * @param hash The hash value we wish to store. - * @param offset Output parameter receiving the chosen file offset index. - * @param evicted_hash Output parameter holding the evicted hash (empty if none). - * @return 1 if newly inserted, 0 if already existed, -1 if temporarily failed (all slots busy) - */ - int acquire_slot(const std::string &hash, size_t &offset, std::string &evicted_hash) { - std::lock_guard lock(index_lock_); - - auto existing = index_.find(hash); - if (existing != index_.end()) { - if (existing->second.writing) { - return -1; // Write in progress, caller should retry - } - lru_list_.splice(lru_list_.begin(), lru_list_, existing->second.lru_iterator); - offset = existing->second.foffset; - evicted_hash.clear(); - return 0; // Already exists and ready - } - - if (!empty_block_list_.empty()) { - offset = empty_block_list_.back(); - empty_block_list_.pop_back(); - evicted_hash.clear(); - } else { - // Need to evict a victim from LRU list - // Find a victim that is not currently being written (writing=false) - auto it = lru_list_.end(); - bool found_victim = false; - - while (it != lru_list_.begin()) { - --it; - const std::string &candidate_hash = *it; - auto candidate_it = index_.find(candidate_hash); - - // Data structure inconsistency detected, return failure to avoid corruption - if (candidate_it == index_.end()) { - fprintf(stderr, "[light_mem error] LRU list and index inconsistent, returning failure\n"); - return -1; - } - - if (!candidate_it->second.writing) { - // Found a valid victim (not currently being written to disk) - offset = candidate_it->second.foffset; - evicted_hash = candidate_hash; - lru_list_.erase(it); - index_.erase(candidate_it); - found_victim = true; - break; - } - } - - if (!found_victim) { - // All slots are busy writing or LRU list is empty - // This is temporary congestion, return -1 to let caller retry - return -1; - } - } - - lru_list_.push_front(hash); - index_[hash] = {lru_list_.begin(), offset, false, true}; // Not ready, writing in progress - return 1; // Newly inserted - } - - /** - * @brief Gets the offset of a hash value - * - * @param hash The hash value to search for - * @return The offset if exists and ready, -1 otherwise - */ - size_t get_offset(const std::string &hash) { - std::lock_guard lock(index_lock_); - auto it = index_.find(hash); - if (it == index_.end() || !it->second.ready || it->second.writing) { - return -1; - } else { - return it->second.foffset; - } - } - - /** - * @brief Mark a hash as ready (data written to disk and readable) - * @return true if the slot is still valid, false if it was evicted - */ - bool mark_ready(const std::string &hash) { - std::lock_guard lock(index_lock_); - auto it = index_.find(hash); - if (it != index_.end()) { - it->second.ready = true; - it->second.writing = false; - return true; - } - return false; // Slot was evicted during write - } - - /** - * @brief Remove a hash entry and recycle its slot - * Used for cleaning up failed writes to prevent zombie slots - */ - void remove(const std::string &hash) { - std::lock_guard lock(index_lock_); - auto it = index_.find(hash); - if (it != index_.end()) { - size_t offset = it->second.foffset; - lru_list_.erase(it->second.lru_iterator); - index_.erase(it); - empty_block_list_.push_back(offset); - } - } - -private: - /** - * @brief Internal structure to store LRU list iterator and exists call count - */ - struct IndexEntry { - std::list::iterator lru_iterator; - size_t foffset; // File pointer - bool ready; // Data is written to disk and readable - bool writing; // Slot is allocated but disk write in progress (evictable but not readable) - }; - - size_t capacity_; ///< Maximum number of hash values to store - std::list lru_list_; ///< LRU list, head is most recent, tail is least recently used - std::list empty_block_list_; ///< List of free disk blocks - std::unordered_map index_; ///< Map from hash value to IndexEntry - std::mutex index_lock_; ///< Mutex protecting index data structures -}; - -class LocalStorageEngine : public StorageEngine { -public: - struct HashInfo { - std::vector> caches; - std::vector> io_locks; - - HashInfo() = default; - }; - - /** - * @brief Constructor for IO sharding. - * @param filename Base filename; actual files will be filename_0, filename_1, ... - * @param storage_size Total file size across all shards. - * @param shard Number of shards. - */ - LocalStorageEngine(const std::string &filename, const size_t storage_size, const size_t shard, - const size_t block_size) - : filename_(filename), storage_size_(storage_size), shard_(shard), block_size_(block_size) { - - // 每个 shard 分到的文件大小 - size_t shard_storage_size = storage_size_ / shard_; - // 每个 shard 能存储的块数 - size_t shard_capacity = shard_storage_size / block_size; - - // 初始化缓存索引、锁和文件对象 - caches_.resize(shard_); - io_locks_.resize(shard_); - files_.resize(shard_); - file_fds_.resize(shard_, -1); - - try { - for (size_t i = 0; i < shard_; i++) { - caches_[i] = std::make_shared(shard_capacity); - io_locks_[i] = std::make_shared(); - } - createOrOpenFiles(shard_storage_size); - } catch (...) { - // Clean up any partially opened files on exception - cleanup(); - throw; - } - } - - ~LocalStorageEngine() override { cleanup(); } - - bool query(const std::string &hash) override { - size_t shard_id = getShard(hash); - return caches_[shard_id]->exists(hash); - } - - size_t write(const char *buf, const std::string &hash) override { - size_t shard_id = getShard(hash); - size_t slot = 0; - std::string evicted_hash; - int result = caches_[shard_id]->acquire_slot(hash, slot, evicted_hash); - (void)evicted_hash; - - // result: 1=newly inserted, 0=already exists, -1=temporarily failed - if (result <= 0) { - return 0; - } - - // 执行磁盘写入，持有 writing=true 防止槽位被驱逐 - size_t offset = slot * block_size_; - try { - std::lock_guard lock(*io_locks_[shard_id]); - files_[shard_id].seekp(offset, std::ios::beg); - files_[shard_id].write(buf, block_size_); - files_[shard_id].flush(); - - // 写入完成后立即标记为 ready，同时释放 writing 标志 - caches_[shard_id]->mark_ready(hash); - } catch (...) { - // 写入失败，清理槽位防止僵尸状态，返回 0 表示失败（不抛异常） - caches_[shard_id]->remove(hash); - return 0; // Write failed, let caller retry - } - - // 立即丢弃新写入的页缓存，避免污染读缓存 - // Note: posix_fadvise is not available on macOS -#ifndef __APPLE__ - if (file_fds_[shard_id] >= 0) { - posix_fadvise(file_fds_[shard_id], offset, block_size_, POSIX_FADV_DONTNEED); - } -#endif - - return block_size_; - } - - size_t read(char *buf, const std::string &hash) override { - size_t shard_id = getShard(hash); - std::lock_guard lock(*io_locks_[shard_id]); - size_t block_idx = caches_[shard_id]->get_offset(hash); - if (block_idx == static_cast(-1)) { - return 0; // Hash does not exist or was evicted (cache miss) - } - - size_t offset = block_idx * block_size_; - try { - files_[shard_id].seekg(offset, std::ios::beg); - files_[shard_id].read(buf, block_size_); - return block_size_; - } catch (...) { - // I/O error during read, return 0 (read failed, treat as cache miss) - fprintf(stderr, "[light_mem error] read: I/O error for hash %s at offset %zu\n", hash.c_str(), offset); - return 0; - } - } - - std::shared_ptr getHashInfo() { - auto info = std::make_shared(); - info->caches = caches_; - info->io_locks = io_locks_; - return info; - } - - bool setHashInfo(const std::shared_ptr &info) { - if (!info) { - fprintf(stderr, "[light_mem error] setHashInfo: HashInfo is null\n"); - return false; - } - if (info->caches.size() != shard_ || info->io_locks.size() != shard_) { - fprintf(stderr, "[light_mem error] setHashInfo: shard size mismatch (expected %zu, got caches=%zu io_locks=%zu)\n", - shard_, info->caches.size(), info->io_locks.size()); - return false; - } - caches_ = info->caches; - io_locks_ = info->io_locks; - return true; - } - -private: - inline size_t getShard(const std::string &hash) { return std::hash{}(hash) % shard_; } - - // Helper function to clean up file resources - void cleanup() { - for (size_t i = 0; i < shard_; i++) { - if (files_[i].is_open()) { - try { - files_[i].close(); - } catch (const std::exception &e) { - fprintf(stderr, "[light_mem warning] LocalStorageEngine::cleanup: failed to close file shard %zu: %s\n", i, - e.what()); - } catch (...) { - fprintf(stderr, - "[light_mem warning] LocalStorageEngine::cleanup: failed to close file shard %zu: unknown error\n", - i); - } - } - if (file_fds_[i] >= 0) { - close(file_fds_[i]); - file_fds_[i] = -1; - } - } - } - - void createOrOpenFiles(size_t shard_storage_size) { - for (size_t i = 0; i < shard_; i++) { - std::stringstream ss; - ss << filename_ << "_" << i; - std::string shard_filename = ss.str(); - - // 打开 fstream 用于读写 - files_[i].open(shard_filename, std::ios::binary | std::ios::in | std::ios::out | std::ios::trunc); - if (!files_[i].is_open()) { - throw std::runtime_error("Failed to open file: " + shard_filename); - } - files_[i].seekp(shard_storage_size - 1, std::ios::beg); - files_[i].write("", 1); - files_[i].seekp(0, std::ios::beg); - - // 同时打开原生 fd 用于 posix_fadvise - file_fds_[i] = open(shard_filename.c_str(), O_RDWR); - if (file_fds_[i] < 0) { - throw std::runtime_error("Failed to open native fd for: " + shard_filename); - } - } - } - - std::string filename_; - size_t storage_size_; - size_t shard_; - size_t block_size_; - std::vector files_; - std::vector file_fds_; ///< 原生文件描述符，用于 posix_fadvise - std::vector> io_locks_; ///< Per-shard mutexes protecting file I/O operations - std::vector> caches_; -}; - -} // namespace storage -} // namespace cache +#pragma once + +#include "storage/local_cache_index.h" +#include "storage/local_storage_wal.h" +#include "storage/redis_client.h" +#include "storage/storage_engine.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include // off_t + +namespace cache { +namespace storage { + +class LocalStorageEngine : public StorageEngine { +public: + struct HashInfo { + std::vector> caches; + std::vector> io_locks; + + HashInfo() = default; + }; + + struct JournalTask { + size_t shard_id; + uint64_t epoch; + uint64_t write_offset; + uint32_t write_len; + uint32_t data_crc; + size_t slot_id; + std::string hash; + std::string evicted_hash; + + std::mutex mu; + std::condition_variable cv; + bool done = false; + bool success = false; + }; + + LocalStorageEngine(const std::string &filename, size_t storage_size, size_t shard, size_t block_size, + const std::string &index_endpoint, const std::string &index_prefix = std::string()); + ~LocalStorageEngine() override; + + bool query(const std::string &hash) override; + + // Batch query variant for high-throughput callers. + // Returns one bool per hash (same order). In online mode this answers "readable on this node". + std::vector queryMany(const std::vector &hashs); + size_t write(const char *buf, const std::string &hash) override; + size_t read(char *buf, const std::string &hash) override; + + // Distributed/online mode control-plane API. + // Update per-shard ownership state as seen by this node. + // - `shard_ids`, `epochs`, `draining` must have the same length (1:1 correspondence). + // - Shards not listed are treated as NOT writable by default. + // - `draining=1` means the shard is in handoff/drain state: still owned/readable, but new writes are rejected. + // - `epoch` is a generation number used to fence writes across ownership changes. + void updateShardAssignments(const std::vector &shard_ids, const std::vector &epochs, + const std::vector &draining); + + // For ownership handoff: rebuild Redis index from local snapshot + WAL for a shard. + // Redis is not a source of truth; this restores the index to match local durable state. + // Designed to be called when a node (re-)acquires write permission for a shard. + void recoverShardToRedis(size_t shard_id); + + // For ownership handoff: choose incremental vs full Redis recovery based on Redis health. + // This is a best-effort optimization; correctness is still guarded by CRC checks. + void recoverShardToRedisSmart(size_t shard_id); + + // Observability: number of in-flight write operations targeting a shard. + // This is best-effort and intended for control-plane decisions (e.g., waiting for draining). + uint32_t shardInflight(size_t shard_id) const; + uint64_t shardWrittenBytes(size_t shard_id) const; + + // True LRU eviction observability (local disk cache). These counts are local to this node. + // `shardEvictionCount`: evictions for a single shard. + uint64_t shardEvictionCount(size_t shard_id) const; + // `evictionCount`: total evictions across all shards. + uint64_t evictionCount() const; + // `evictionObserved`: quick boolean check (evictionCount() > 0). + bool evictionObserved() const; + + std::shared_ptr getHashInfo(); + bool setHashInfo(const std::shared_ptr &info); + +private: + size_t getShard(const std::string &hash) const; + bool isShardWritable(size_t shard_id, uint64_t *epoch_out = nullptr) const; + std::optional findShardInRedis(const std::string &hash, size_t *slot_id_out); + size_t pickWritableShard(const std::string &hash) const; + + // Best-effort local hint to avoid O(shards) scans and repeated Redis lookups. + // Only used in online_mode_ and always re-validated under the shard io lock. + std::optional localShardHint(const std::string &hash) const; + void noteLocalShardHint(const std::string &hash, size_t shard_id); + void eraseLocalShardHint(const std::string &hash); + + void onEpochChangedLocked(size_t shard_id, uint64_t new_epoch); + void appendEpochMarkerLocked(size_t shard_id, uint64_t epoch); + + void startJournalWorkers(); + void stopJournalWorkers(); + void journalWorkerLoop(size_t shard_id); + + void cleanup(); + void createOrOpenFiles(size_t shard_storage_size); + void initIndexBackend(const std::string &endpoint, const std::string &index_prefix); + + bool preadAll(int fd, void *buf, size_t len, off_t offset); + bool pwriteAll(int fd, const void *buf, size_t len, off_t offset); + + bool readSuperBlockAt(size_t shard_id, off_t off, SuperBlock &sb); + void writeSuperBlockAt(size_t shard_id, off_t off, SuperBlock sb); + void ensureSuperBlocksInitialized(size_t shard_id); + + void appendJournalRecord(size_t shard_id, uint64_t write_offset, uint32_t write_len, uint32_t data_crc, + const std::string &hash, const std::string &evicted_hash); + + void maybeCheckpoint(size_t shard_id); + void checkpoint(size_t shard_id); + void writeCheckpointSuperblockOnly(size_t shard_id); + + void truncateJournalToHeader(size_t shard_id); + + struct JournalOp { + std::string hash; + std::string evicted; + size_t slot_id = 0; + uint32_t data_crc = 0; + }; + void scanWalOps(size_t shard_id, size_t shard_capacity, off_t start_off, std::vector &ops); + + void recoverAllShards(size_t shard_capacity); + void recoverShard(size_t shard_id, size_t shard_capacity); + void recoverShardToRedisIncremental(size_t shard_id); + bool shouldFullRecoverRedis(size_t shard_id); + + std::string filename_; + size_t storage_size_; + size_t shard_; + size_t block_size_; + + std::vector file_fds_; + std::vector meta_fds_; + + std::vector> io_locks_; + std::vector> caches_; + + // Redis is used for: + // - global dedupe/locking on the write hot path + // - publishing shard/global index updates from journal workers + // A single connection becomes a bottleneck under high concurrency; keep a + // dedicated connection for locks and a small pool for shard/journal updates. + std::unique_ptr redis_lock_; + std::vector> redis_pool_; + std::vector journal_entries_; + std::vector superblock_seq_; + std::vector superblock_stable_offset_; + std::vector superblock_epoch_; + std::vector shard_recovered_epoch_; + + // Shard write permission cache (fed by coordinator). + // NOTE: std::vector> is ill-formed on libstdc++ because atomic is non-movable. + // Use fixed-size arrays instead. + std::unique_ptr[]> shard_writable_; + std::unique_ptr[]> shard_draining_; + std::unique_ptr[]> shard_epoch_cache_; + std::unique_ptr[]> shard_inflight_; + std::unique_ptr[]> shard_written_bytes_; + + std::vector journal_threads_; + std::vector> journal_mu_; + std::vector> journal_cv_; + std::vector>> journal_queue_; + std::vector journal_stop_; + + // When online mode is enabled, shards can be dynamically assigned and hash->shard is not deterministic. + // In the default single-node mode, keep deterministic sharding to make query/write O(1) per hash. + bool online_mode_ = false; + + // hash -> shard_id (bounded by eviction). + // NOTE: This is on the hot read path in online mode. Using a single global shared_mutex + // can become a scalability bottleneck (many read threads contending on one lock). + // Shard it into multiple buckets to reduce contention. + static constexpr size_t kLocalHintBuckets = 64; + static_assert((kLocalHintBuckets & (kLocalHintBuckets - 1)) == 0, "kLocalHintBuckets must be power-of-two"); + + size_t localHintBucket(const std::string &hash) const { + return std::hash{}(hash) & (kLocalHintBuckets - 1); + } + + mutable std::array local_hint_mu_; + std::array, kLocalHintBuckets> local_hash_to_shard_; + + RedisClient *redisForShard(size_t shard_id) const { + if (redis_pool_.empty()) { + return nullptr; + } + return redis_pool_[shard_id % redis_pool_.size()].get(); + } +}; + +} // namespace storage +} // namespace cache diff --git a/src/storage/local_storage_engine_files.cpp b/src/storage/local_storage_engine_files.cpp new file mode 100755 index 0000000..a4e46ba --- /dev/null +++ b/src/storage/local_storage_engine_files.cpp @@ -0,0 +1,400 @@ +#include "storage/local_storage_engine.h" + +#include "config.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace cache { +namespace storage { + +static void chmod_best_effort(const std::string &path, mode_t mode) { + // Best-effort: on shared filesystems the owner should be able to chmod. + // If it fails (e.g., permission/readonly), leave it as-is. + (void)::chmod(path.c_str(), mode); +} + +static void ensure_parent_dir_shareable(const std::string &path) { + namespace fs = std::filesystem; + std::error_code ec; + + fs::path p(path); + fs::path parent = p.parent_path(); + if (parent.empty()) { + parent = fs::current_path(ec); + if (ec) { + // If current_path fails, just skip parent handling. + return; + } + } + + ec.clear(); + fs::create_directories(parent, ec); + // Even if create_directories fails (e.g. already exists), try chmod. + chmod_best_effort(parent.string(), 0777); +} + +static void ensure_directory(const std::string &path) { + struct stat st; + if (::stat(path.c_str(), &st) != 0) { + if (::mkdir(path.c_str(), 0777) != 0 && errno != EEXIST) { + const int err = errno; + throw std::runtime_error("Failed to create directory: " + path + ", errno=" + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); + } + } else if (!S_ISDIR(st.st_mode)) { + throw std::runtime_error("Path exists but is not a directory: " + path); + } + chmod_best_effort(path, 0777); +} + +static void require_directory_exists(const std::string &path) { + struct stat st; + if (::stat(path.c_str(), &st) != 0) { + const int err = errno; + throw std::runtime_error("Missing directory: " + path + ", errno=" + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); + } + if (!S_ISDIR(st.st_mode)) { + throw std::runtime_error("Path exists but is not a directory: " + path); + } +} + +// Initializer-only: open existing OR create new file without truncating an existing file. +// If created, optionally preallocates and writes an initial header. +static int open_existing_or_create_new(const std::string &path, size_t preallocate_size, const void *init_data, + size_t init_size) { + bool created = false; + int fd = ::open(path.c_str(), O_RDWR | O_CREAT | O_EXCL, 0666); + if (fd >= 0) { + created = true; + } else if (errno == EEXIST) { + fd = ::open(path.c_str(), O_RDWR); + } + if (fd < 0) { + const int err = errno; + throw std::runtime_error("Failed to open file: " + path + ", errno=" + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); + } + + if (created) { + if (preallocate_size > 0) { + if (::lseek(fd, static_cast(preallocate_size - 1), SEEK_SET) < 0 || ::write(fd, "", 1) != 1) { + const int err = errno; + ::close(fd); + throw std::runtime_error("Failed to preallocate file: " + path + ", errno=" + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); + } + } + if (init_data && init_size > 0) { + if (::pwrite(fd, init_data, init_size, 0) != static_cast(init_size)) { + const int err = errno; + ::close(fd); + throw std::runtime_error("Failed to initialize file: " + path + ", errno=" + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); + } + } + } + + chmod_best_effort(path, 0666); + return fd; +} + +// Follower-only: open existing file; never creates/truncates. +static int open_existing_file(const std::string &path) { + int fd = ::open(path.c_str(), O_RDWR); + if (fd < 0) { + const int err = errno; + throw std::runtime_error("Failed to open existing file: " + path + ", errno=" + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); + } + return fd; +} + +static bool file_exists(const std::string &path) { + struct stat st; + return (::stat(path.c_str(), &st) == 0); +} + +static void write_init_marker(const std::string &path) { + int fd = ::open(path.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0666); + if (fd < 0) { + const int err = errno; + throw std::runtime_error("Failed to write init marker: " + path + ", errno=" + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); + } + static const char kMsg[] = "ok\n"; + (void)::write(fd, kMsg, sizeof(kMsg) - 1); + (void)::fsync(fd); + ::close(fd); + chmod_best_effort(path, 0666); +} + +void LocalStorageEngine::cleanup() { + for (size_t i = 0; i < shard_; i++) { + if (file_fds_[i] >= 0) { + ::close(file_fds_[i]); + file_fds_[i] = -1; + } + if (meta_fds_[i] >= 0) { + ::close(meta_fds_[i]); + meta_fds_[i] = -1; + } + } +} + +void LocalStorageEngine::createOrOpenFiles(size_t shard_storage_size) { + ensure_parent_dir_shareable(filename_); + + static const char kMetaHeader[META_HEADER_SIZE] = {0}; + + // Global init lock to avoid TOCTOU races when multiple services start against the same directory. + // Only the initializer creates directories/files; followers wait for the marker and then only open existing files. + const std::string init_lock_dir = filename_ + ".init.lock"; + const std::string init_marker = filename_ + ".initialized"; + + bool initializer = false; + if (::mkdir(init_lock_dir.c_str(), 0777) == 0) { + initializer = true; + chmod_best_effort(init_lock_dir, 0777); + } else if (errno == EEXIST) { + initializer = false; + } else { + const int err = errno; + throw std::runtime_error("Failed to create init lock: " + init_lock_dir + ", errno=" + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); + } + + if (!initializer) { + // If marker already exists, proceed immediately. Otherwise, wait for initializer to finish. + if (!file_exists(init_marker)) { + static constexpr int kWaitMs = 50; + static constexpr int kTimeoutMs = 60000; // 60s + int waited = 0; + while (!file_exists(init_marker) && waited < kTimeoutMs) { + ::usleep(kWaitMs * 1000); + waited += kWaitMs; + } + if (!file_exists(init_marker)) { + throw std::runtime_error("Timed out waiting for storage init marker: " + init_marker + + ". Another process may be stuck initializing (lock=" + init_lock_dir + ")"); + } + } + } + + for (size_t i = 0; i < shard_; i++) { + const std::string shard_path = filename_ + "_" + std::to_string(i); + if (initializer) { + ensure_directory(shard_path); + } else { + require_directory_exists(shard_path); + } + + const std::string data_filename = shard_path + "/data"; + const std::string meta_filename = shard_path + "/meta"; + + if (initializer) { + // Open data file (preallocate only when newly created). + file_fds_[i] = open_existing_or_create_new(data_filename, shard_storage_size, nullptr, 0); + // Open meta file (initialize header only when newly created). + meta_fds_[i] = open_existing_or_create_new(meta_filename, 0, kMetaHeader, META_HEADER_SIZE); + } else { + // Follower: never create/truncate. + file_fds_[i] = open_existing_file(data_filename); + meta_fds_[i] = open_existing_file(meta_filename); + } + + ensureSuperBlocksInitialized(i); + } + + if (initializer) { + // Publish completion marker then release lock. + write_init_marker(init_marker); + (void)::rmdir(init_lock_dir.c_str()); + } +} + +void LocalStorageEngine::initIndexBackend(const std::string &endpoint, const std::string &index_prefix) { + // Index backend is opt-in: only enable when endpoint is set. + std::string ep = endpoint; + if (ep.empty()) { + return; + } + + RedisClient::Options opt; + + // Optional: override Redis key prefix to isolate multiple runs. + // Unified with coordinator prefix via the single init arg `index_prefix`. + { + std::string v = index_prefix; + while (!v.empty() && (v.front() == ' ' || v.front() == '\t')) { + v.erase(v.begin()); + } + while (!v.empty() && (v.back() == ' ' || v.back() == '\t')) { + v.pop_back(); + } + if (!v.empty()) { + opt.key_prefix = std::move(v); + } + } + + // Strip scheme if someone passes it. + const std::string http_prefix = "http://"; + const std::string https_prefix = "https://"; + if (ep.rfind(http_prefix, 0) == 0) { + ep = ep.substr(http_prefix.size()); + } else if (ep.rfind(https_prefix, 0) == 0) { + ep = ep.substr(https_prefix.size()); + } + + // Take first if comma-separated. + const auto comma = ep.find(','); + if (comma != std::string::npos) { + ep = ep.substr(0, comma); + } + + // Trim whitespace (minimal). + while (!ep.empty() && (ep.front() == ' ' || ep.front() == '\t')) { + ep.erase(ep.begin()); + } + while (!ep.empty() && (ep.back() == ' ' || ep.back() == '\t')) { + ep.pop_back(); + } + + if (!ep.empty()) { + auto colon = ep.rfind(':'); + if (colon != std::string::npos) { + const std::string host_part = ep.substr(0, colon); + const std::string port_part = ep.substr(colon + 1); + if (!host_part.empty()) { + opt.host = host_part; + } + if (!port_part.empty()) { + try { + opt.port = std::stoi(port_part); + } catch (...) { + } + } + } else { + opt.host = ep; + } + } + + // Pool size: controls number of TCP connections used for shard/journal metadata updates. + // A larger pool reduces mutex/queue contention under high concurrency. + int pool_size = 32; + + // Dedicated client for write-hot-path lock/dedupe. + redis_lock_ = std::make_unique(opt); + + // Pool for journal workers / lookups. + redis_pool_.clear(); + redis_pool_.reserve(static_cast(pool_size)); + for (int i = 0; i < pool_size; i++) { + redis_pool_.emplace_back(std::make_unique(opt)); + } +} + +bool LocalStorageEngine::preadAll(int fd, void *buf, size_t len, off_t offset) { + char *p = static_cast(buf); + size_t left = len; + while (left > 0) { + ssize_t n = ::pread(fd, p, left, offset); + if (n <= 0) { + return false; + } + p += static_cast(n); + left -= static_cast(n); + offset += static_cast(n); + } + return true; +} + +bool LocalStorageEngine::pwriteAll(int fd, const void *buf, size_t len, off_t offset) { + const char *p = static_cast(buf); + size_t left = len; + while (left > 0) { + ssize_t n = ::pwrite(fd, p, left, offset); + if (n <= 0) { + return false; + } + p += static_cast(n); + left -= static_cast(n); + offset += static_cast(n); + } + return true; +} + +bool LocalStorageEngine::readSuperBlockAt(size_t shard_id, off_t off, SuperBlock &sb) { + if (!preadAll(meta_fds_[shard_id], &sb, sizeof(SuperBlock), off)) { + return false; + } + if (sb.magic != SUPERBLOCK_MAGIC || sb.version != SUPERBLOCK_VERSION || sb.shard_id != shard_id) { + return false; + } + const uint32_t expect = compute_crc32(&sb, SUPERBLOCK_SIZE - sizeof(uint32_t)); + if (expect != sb.crc32) { + return false; + } + if (sb.block_size != block_size_) { + return false; + } + return true; +} + +void LocalStorageEngine::writeSuperBlockAt(size_t shard_id, off_t off, SuperBlock sb) { + sb.crc32 = 0; + sb.crc32 = compute_crc32(&sb, SUPERBLOCK_SIZE - sizeof(uint32_t)); + if (!pwriteAll(meta_fds_[shard_id], &sb, sizeof(SuperBlock), off)) { + throw std::runtime_error("Failed to write superblock"); + } + ::fsync(meta_fds_[shard_id]); +} + +void LocalStorageEngine::ensureSuperBlocksInitialized(size_t shard_id) { + SuperBlock a{}, b{}; + bool va = readSuperBlockAt(shard_id, 0, a); + bool vb = readSuperBlockAt(shard_id, static_cast(SUPERBLOCK_SIZE), b); + + if (va || vb) { + const SuperBlock &chosen = (!va) ? b : (!vb) ? a : (a.sequence_id >= b.sequence_id ? a : b); + superblock_seq_[shard_id] = chosen.sequence_id; + superblock_stable_offset_[shard_id] = + (chosen.stable_offset >= META_HEADER_SIZE) ? chosen.stable_offset : META_HEADER_SIZE; + superblock_epoch_[shard_id] = chosen.current_epoch; + return; + } + + SuperBlock sb{}; + sb.magic = SUPERBLOCK_MAGIC; + sb.version = SUPERBLOCK_VERSION; + sb.shard_id = shard_id; + sb.sequence_id = 1; + sb.stable_offset = META_HEADER_SIZE; + sb.total_capacity = (storage_size_ / shard_) / block_size_; + sb.block_size = static_cast(block_size_); + sb.current_epoch = 0; + writeSuperBlockAt(shard_id, 0, sb); + sb.sequence_id = 0; + writeSuperBlockAt(shard_id, static_cast(SUPERBLOCK_SIZE), sb); + superblock_seq_[shard_id] = 1; + superblock_stable_offset_[shard_id] = META_HEADER_SIZE; + superblock_epoch_[shard_id] = 0; + + if (::ftruncate(meta_fds_[shard_id], META_HEADER_SIZE) != 0) { + throw std::runtime_error("Failed to ftruncate meta to header"); + } + (void)::fsync(meta_fds_[shard_id]); +} + +} // namespace storage +} // namespace cache diff --git a/src/storage/local_storage_engine_journal.cpp b/src/storage/local_storage_engine_journal.cpp new file mode 100755 index 0000000..eb1a4c8 --- /dev/null +++ b/src/storage/local_storage_engine_journal.cpp @@ -0,0 +1,334 @@ +#include "storage/local_storage_engine.h" + +#include + +#include "utils/fsync_compat.h" + +#include +#include +#include +#include + +namespace cache { +namespace storage { + +void LocalStorageEngine::startJournalWorkers() { + journal_threads_.resize(shard_); + for (size_t shard_id = 0; shard_id < shard_; shard_id++) { + journal_threads_[shard_id] = std::thread([this, shard_id] { journalWorkerLoop(shard_id); }); + } +} + +void LocalStorageEngine::stopJournalWorkers() { + for (size_t shard_id = 0; shard_id < shard_; shard_id++) { + { + std::lock_guard lk(*journal_mu_[shard_id]); + journal_stop_[shard_id] = true; + } + journal_cv_[shard_id]->notify_one(); + } + for (size_t shard_id = 0; shard_id < shard_; shard_id++) { + if (journal_threads_.size() > shard_id && journal_threads_[shard_id].joinable()) { + journal_threads_[shard_id].join(); + } + } +} + +void LocalStorageEngine::journalWorkerLoop(size_t shard_id) { + while (true) { + std::vector> batch; + { + std::unique_lock lk(*journal_mu_[shard_id]); + journal_cv_[shard_id]->wait(lk, [&] { return journal_stop_[shard_id] || !journal_queue_[shard_id].empty(); }); + if (journal_queue_[shard_id].empty()) { + if (journal_stop_[shard_id]) { + break; + } + continue; + } + while (!journal_queue_[shard_id].empty()) { + batch.emplace_back(journal_queue_[shard_id].front()); + journal_queue_[shard_id].pop_front(); + } + } + + bool ok = true; + try { + { + std::unique_lock io_lock(*io_locks_[shard_id]); + + for (auto &task : batch) { + if (task->epoch != 0 && task->epoch != superblock_epoch_[shard_id]) { + onEpochChangedLocked(shard_id, task->epoch); + } + appendJournalRecord(shard_id, task->write_offset, task->write_len, task->data_crc, task->hash, + task->evicted_hash); + journal_entries_[shard_id]++; + } + + writeCheckpointSuperblockOnly(shard_id); + } + + RedisClient *r = redisForShard(shard_id); + const bool redis_connected = (r && r->connect()); + if (redis_connected) { + const std::string key = r->shardIndexKey(shard_id); + const std::string gkey = r->globalIndexKey(); + + std::vector> cmds; + cmds.reserve(batch.size() * 4 + 1); + + for (auto &task : batch) { + if (!task->evicted_hash.empty()) { + cmds.push_back({"HDEL", key, task->evicted_hash}); + cmds.push_back({"HDEL", gkey, task->evicted_hash}); + cmds.push_back({"HDEL", r->globalCrcKey(), task->evicted_hash}); + } + cmds.push_back({"HSET", key, task->hash, std::to_string(task->slot_id)}); + cmds.push_back({"HSET", gkey, task->hash, std::to_string(shard_id) + ":" + std::to_string(task->slot_id)}); + cmds.push_back({"HSET", r->globalCrcKey(), task->hash, std::to_string(task->data_crc)}); + } + cmds.push_back({"SET", r->shardSeqKey(shard_id), std::to_string(superblock_seq_[shard_id])}); + + if (!r->pipeline(cmds)) { + std::fprintf(stderr, "[light_mem error] journalWorkerLoop: Redis pipeline failed (shard=%zu, cmds=%zu)\n", + shard_id, cmds.size()); + ok = false; + } + } + + // Keep local WAL bounded even when Redis is down. + // Only checkpoint after the write queue is drained (to avoid frequent checkpointing + // in the middle of a steady stream of writes). + bool queue_empty = false; + { + std::lock_guard lk(*journal_mu_[shard_id]); + queue_empty = journal_queue_[shard_id].empty(); + } + if (queue_empty) { + maybeCheckpoint(shard_id); + } + } catch (...) { + ok = false; + } + + for (auto &task : batch) { + { + std::lock_guard lk(task->mu); + task->success = ok; + task->done = true; + } + task->cv.notify_one(); + } + } + + // Drain any queued tasks (fail them) if shutting down. + while (true) { + std::shared_ptr task; + { + std::lock_guard lk(*journal_mu_[shard_id]); + if (journal_queue_[shard_id].empty()) { + break; + } + task = journal_queue_[shard_id].front(); + journal_queue_[shard_id].pop_front(); + } + { + std::lock_guard lk(task->mu); + task->success = false; + task->done = true; + } + task->cv.notify_one(); + } +} + +void LocalStorageEngine::appendJournalRecord(size_t shard_id, uint64_t write_offset, uint32_t write_len, + uint32_t data_crc, const std::string &hash, + const std::string &evicted_hash) { + JournalRecord rec{}; + rec.magic = JOURNAL_MAGIC; + rec.version = JOURNAL_VERSION; + rec.flags = 0; + rec.epoch = superblock_epoch_[shard_id]; + rec.write_offset = write_offset; + rec.write_len = write_len; + rec.data_crc = data_crc; + rec.hash_len = static_cast(hash.size()); + rec.evicted_hash_len = static_cast(evicted_hash.size()); + + std::vector tmp; + tmp.resize(sizeof(JournalRecord) + hash.size() + evicted_hash.size()); + std::memcpy(tmp.data(), &rec, sizeof(JournalRecord)); + if (!hash.empty()) { + std::memcpy(tmp.data() + sizeof(JournalRecord), hash.data(), hash.size()); + } + if (!evicted_hash.empty()) { + std::memcpy(tmp.data() + sizeof(JournalRecord) + hash.size(), evicted_hash.data(), evicted_hash.size()); + } + const uint32_t record_crc = compute_crc32(tmp.data(), tmp.size()); + + off_t end = ::lseek(meta_fds_[shard_id], 0, SEEK_END); + if (end < 0) { + throw std::runtime_error("Failed to seek meta file end"); + } + + if (!pwriteAll(meta_fds_[shard_id], &rec, sizeof(JournalRecord), end)) { + throw std::runtime_error("Failed to append journal header"); + } + end += static_cast(sizeof(JournalRecord)); + + if (!hash.empty()) { + if (!pwriteAll(meta_fds_[shard_id], hash.data(), hash.size(), end)) { + throw std::runtime_error("Failed to append journal hash"); + } + end += static_cast(hash.size()); + } + + if (!evicted_hash.empty()) { + if (!pwriteAll(meta_fds_[shard_id], evicted_hash.data(), evicted_hash.size(), end)) { + throw std::runtime_error("Failed to append journal evicted_hash"); + } + end += static_cast(evicted_hash.size()); + } + + if (!pwriteAll(meta_fds_[shard_id], &record_crc, sizeof(uint32_t), end)) { + throw std::runtime_error("Failed to append journal crc"); + } + if (cache::utils::fdatasync_compat(meta_fds_[shard_id]) != 0) { + throw std::runtime_error("Failed to fdatasync meta"); + } +} + +void LocalStorageEngine::appendEpochMarkerLocked(size_t shard_id, uint64_t epoch) { + JournalRecord rec{}; + rec.magic = JOURNAL_MAGIC; + rec.version = JOURNAL_VERSION; + rec.flags = JOURNAL_FLAG_EPOCH_MARKER; + rec.epoch = epoch; + rec.write_offset = 0; + rec.write_len = 0; + rec.data_crc = 0; + rec.hash_len = 0; + rec.evicted_hash_len = 0; + + std::vector tmp; + tmp.resize(sizeof(JournalRecord)); + std::memcpy(tmp.data(), &rec, sizeof(JournalRecord)); + const uint32_t record_crc = compute_crc32(tmp.data(), tmp.size()); + + off_t end = ::lseek(meta_fds_[shard_id], 0, SEEK_END); + if (end < 0) { + throw std::runtime_error("Failed to seek meta file end (epoch marker)"); + } + + if (!pwriteAll(meta_fds_[shard_id], &rec, sizeof(JournalRecord), end)) { + throw std::runtime_error("Failed to append epoch marker"); + } + end += static_cast(sizeof(JournalRecord)); + + if (!pwriteAll(meta_fds_[shard_id], &record_crc, sizeof(uint32_t), end)) { + throw std::runtime_error("Failed to append epoch marker crc"); + } + if (cache::utils::fdatasync_compat(meta_fds_[shard_id]) != 0) { + throw std::runtime_error("Failed to fdatasync meta (epoch marker)"); + } +} + +void LocalStorageEngine::onEpochChangedLocked(size_t shard_id, uint64_t new_epoch) { + // Persist the new epoch into memory and WAL so later scans can ignore stale tail writes. + superblock_epoch_[shard_id] = new_epoch; + shard_epoch_cache_[shard_id].store(new_epoch, std::memory_order_relaxed); + appendEpochMarkerLocked(shard_id, new_epoch); + writeCheckpointSuperblockOnly(shard_id); +} + +void LocalStorageEngine::maybeCheckpoint(size_t shard_id) { + // Trigger checkpoint based on accumulated write volume (about 1GiB), which scales with block size. + // NOTE: block_size_ is derived from the cache service block size and is affected by + static constexpr uint64_t kCheckpointBytes = 1ull * 1024ull * 1024ull * 1024ull; // 1GiB + + const uint64_t bs = static_cast(block_size_); + if (bs == 0) { + return; + } + const uint64_t blocks_per_checkpoint = (kCheckpointBytes + bs - 1) / bs; // ceil + const uint64_t threshold = (blocks_per_checkpoint == 0) ? 1 : blocks_per_checkpoint; + + if (journal_entries_[shard_id] < threshold) { + return; + } + checkpoint(shard_id); +} + +void LocalStorageEngine::checkpoint(size_t shard_id) { + std::unique_lock lock(*io_locks_[shard_id]); + + // 1. Create Local Index Snapshot + std::stringstream ss; + ss << filename_ << "_" << shard_id << "/index"; + std::string snap_path = ss.str(); + + if (!caches_[shard_id]->saveToSnapshot(snap_path)) { + std::fprintf(stderr, "[light_mem warning] checkpoint: failed to save snapshot for shard %zu\n", shard_id); + // We continue even if snapshot fails? No, if snapshot fails, we shouldn't truncate WAL. + return; + } + + SuperBlock sb{}; + sb.magic = SUPERBLOCK_MAGIC; + sb.version = SUPERBLOCK_VERSION; + sb.shard_id = shard_id; + sb.sequence_id = superblock_seq_[shard_id] + 1; + sb.stable_offset = META_HEADER_SIZE; + sb.total_capacity = (storage_size_ / shard_) / block_size_; + sb.block_size = static_cast(block_size_); + sb.current_epoch = superblock_epoch_[shard_id]; + + const off_t target = (sb.sequence_id % 2 == 1) ? 0 : static_cast(SUPERBLOCK_SIZE); + writeSuperBlockAt(shard_id, target, sb); + superblock_seq_[shard_id] = sb.sequence_id; + superblock_stable_offset_[shard_id] = sb.stable_offset; + + // Best-effort: publish local sequence so Redis consumers can detect staleness. + RedisClient *r = redisForShard(shard_id); + if (r && r->connect()) { + (void)r->setString(r->shardSeqKey(shard_id), std::to_string(superblock_seq_[shard_id])); + } + + if (::ftruncate(meta_fds_[shard_id], META_HEADER_SIZE) != 0) { + return; + } + (void)::fsync(meta_fds_[shard_id]); + journal_entries_[shard_id] = 0; +} + +void LocalStorageEngine::writeCheckpointSuperblockOnly(size_t shard_id) { + SuperBlock sb{}; + sb.magic = SUPERBLOCK_MAGIC; + sb.version = SUPERBLOCK_VERSION; + sb.shard_id = shard_id; + sb.sequence_id = superblock_seq_[shard_id] + 1; + sb.stable_offset = META_HEADER_SIZE; + sb.total_capacity = (storage_size_ / shard_) / block_size_; + sb.block_size = static_cast(block_size_); + sb.current_epoch = superblock_epoch_[shard_id]; + + const off_t target = (sb.sequence_id % 2 == 1) ? 0 : static_cast(SUPERBLOCK_SIZE); + writeSuperBlockAt(shard_id, target, sb); + superblock_seq_[shard_id] = sb.sequence_id; + superblock_stable_offset_[shard_id] = sb.stable_offset; +} + +void LocalStorageEngine::truncateJournalToHeader(size_t shard_id) { + if (meta_fds_[shard_id] < 0) { + return; + } + if (::ftruncate(meta_fds_[shard_id], META_HEADER_SIZE) != 0) { + return; + } + (void)::fsync(meta_fds_[shard_id]); + journal_entries_[shard_id] = 0; +} + +} // namespace storage +} // namespace cache diff --git a/src/storage/local_storage_engine_public.cpp b/src/storage/local_storage_engine_public.cpp new file mode 100755 index 0000000..145cbc4 --- /dev/null +++ b/src/storage/local_storage_engine_public.cpp @@ -0,0 +1,615 @@ +#include "storage/local_storage_engine.h" + +#include "config.h" + +#include +#include + +#include "utils/fsync_compat.h" + +#include +#include +#include +#include +#include + +namespace cache { +namespace storage { + +std::optional LocalStorageEngine::localShardHint(const std::string &hash) const { + const size_t b = localHintBucket(hash); + std::shared_lock lk(local_hint_mu_[b]); + auto it = local_hash_to_shard_[b].find(hash); + if (it == local_hash_to_shard_[b].end()) { + return std::nullopt; + } + return it->second; +} + +void LocalStorageEngine::noteLocalShardHint(const std::string &hash, size_t shard_id) { + // Fast-path: avoid exclusive lock when the mapping is already up-to-date. + const size_t b = localHintBucket(hash); + { + std::shared_lock lk(local_hint_mu_[b]); + auto it = local_hash_to_shard_[b].find(hash); + if (it != local_hash_to_shard_[b].end() && it->second == shard_id) { + return; + } + } + std::unique_lock lk(local_hint_mu_[b]); + local_hash_to_shard_[b][hash] = shard_id; +} + +void LocalStorageEngine::eraseLocalShardHint(const std::string &hash) { + // Fast-path: avoid exclusive lock if the key is absent. + const size_t b = localHintBucket(hash); + { + std::shared_lock lk(local_hint_mu_[b]); + if (local_hash_to_shard_[b].find(hash) == local_hash_to_shard_[b].end()) { + return; + } + } + std::unique_lock lk(local_hint_mu_[b]); + local_hash_to_shard_[b].erase(hash); +} + +LocalStorageEngine::LocalStorageEngine(const std::string &filename, const size_t storage_size, const size_t shard, + const size_t block_size, const std::string &index_endpoint, + const std::string &index_prefix) + : filename_(filename), storage_size_(storage_size), shard_(shard), block_size_(block_size), + online_mode_(!index_endpoint.empty()) { + // 每个 shard 分到的文件大小 + const size_t shard_storage_size = storage_size_ / shard_; + // 每个 shard 能存储的块数 + const size_t shard_capacity = shard_storage_size / block_size; + + caches_.resize(shard_); + io_locks_.resize(shard_); + file_fds_.resize(shard_, -1); + meta_fds_.resize(shard_, -1); + journal_entries_.assign(shard_, 0); + superblock_seq_.assign(shard_, 0); + superblock_stable_offset_.assign(shard_, META_HEADER_SIZE); + superblock_epoch_.assign(shard_, 0); + shard_recovered_epoch_.assign(shard_, 0); + + shard_writable_ = std::make_unique[]>(shard_); + shard_draining_ = std::make_unique[]>(shard_); + shard_epoch_cache_ = std::make_unique[]>(shard_); + shard_inflight_ = std::make_unique[]>(shard_); + shard_written_bytes_ = std::make_unique[]>(shard_); + for (size_t i = 0; i < shard_; i++) { + shard_writable_[i].store(1, std::memory_order_relaxed); + shard_draining_[i].store(0, std::memory_order_relaxed); + shard_epoch_cache_[i].store(0, std::memory_order_relaxed); + shard_inflight_[i].store(0, std::memory_order_relaxed); + shard_written_bytes_[i].store(0, std::memory_order_relaxed); + } + + journal_mu_.resize(shard_); + journal_cv_.resize(shard_); + journal_queue_.resize(shard_); + journal_stop_.assign(shard_, false); + + initIndexBackend(index_endpoint, index_prefix); + + // Non-strict by default: if index backend is configured but unreachable, warn and proceed without it. + if (redis_lock_) { + if (!redis_lock_->connect()) { + std::fprintf(stderr, "[light_mem warning] index backend is configured but not reachable; " + "continuing without index persistence\n"); + redis_lock_.reset(); + redis_pool_.clear(); + } + } + + try { + for (size_t i = 0; i < shard_; i++) { + caches_[i] = std::make_shared(shard_capacity); + io_locks_[i] = std::make_shared(); + journal_mu_[i] = std::make_unique(); + journal_cv_[i] = std::make_unique(); + } + createOrOpenFiles(shard_storage_size); + if (!online_mode_) { + recoverAllShards(shard_capacity); + } + startJournalWorkers(); + } catch (...) { + cleanup(); + throw; + } +} + +LocalStorageEngine::~LocalStorageEngine() { + stopJournalWorkers(); + + // Best-effort final checkpoint: + // - Persist an index snapshot + // - Truncate WAL even if it didn't hit the periodic checkpoint threshold + // This prevents startup-time WAL scans from growing unbounded across runs. + for (size_t shard_id = 0; shard_id < shard_; shard_id++) { + if (journal_entries_.size() <= shard_id) { + continue; + } + if (journal_entries_[shard_id] == 0) { + continue; + } + try { + checkpoint(shard_id); + } catch (const std::exception &e) { + std::fprintf(stderr, "[light_mem warning] final checkpoint failed for shard %zu: %s\n", shard_id, e.what()); + } catch (...) { + std::fprintf(stderr, "[light_mem warning] final checkpoint failed for shard %zu: unknown error\n", shard_id); + } + } + + cleanup(); +} + +bool LocalStorageEngine::query(const std::string &hash) { + if (!online_mode_) { + const size_t shard_id = getShard(hash); + return caches_[shard_id]->exists(hash); + } + + // Best-effort hot path: if we already learned hash->shard locally and we still own that shard, + // validate via shard-local index and avoid Redis. + { + auto hinted = localShardHint(hash); + if (hinted.has_value() && hinted.value() < shard_) { + const size_t shard_id = hinted.value(); + if (caches_[shard_id]->exists(hash)) { + return true; + } else { + // Stale hint (evicted locally). + eraseLocalShardHint(hash); + } + } + } + + // Distributed mode: + // - If Redis is configured, prefer the global index to avoid O(shards) scans. + // - If Redis is unavailable, fall back to scanning local in-memory indices. + if (redis_lock_ && redis_lock_->connect()) { + return redis_lock_->hexists(redis_lock_->globalIndexKey(), hash); + } + for (size_t i = 0; i < shard_; i++) { + if (caches_[i]->exists(hash)) { + return true; + } + } + return false; +} + +std::vector LocalStorageEngine::queryMany(const std::vector &hashs) { + std::vector ret; + ret.assign(hashs.size(), false); + + if (!online_mode_) { + for (size_t i = 0; i < hashs.size(); i++) { + const std::string &hash = hashs[i]; + const size_t shard_id = getShard(hash); + ret[i] = caches_[shard_id]->exists(hash); + } + return ret; + } + + // First pass: local hint + shard-local index. + // We only treat it as a hit without Redis re-validation if the shard is currently owned. + std::vector need_redis_hash; + std::vector need_redis_idx; + need_redis_hash.reserve(hashs.size()); + need_redis_idx.reserve(hashs.size()); + + for (size_t i = 0; i < hashs.size(); i++) { + const std::string &hash = hashs[i]; + auto hinted = localShardHint(hash); + if (hinted.has_value() && hinted.value() < shard_) { + const size_t shard_id = hinted.value(); + if (caches_[shard_id]->exists(hash)) { + ret[i] = true; + continue; + } else { + // Stale hint (evicted locally). + eraseLocalShardHint(hash); + } + } + need_redis_hash.emplace_back(hash); + need_redis_idx.emplace_back(i); + } + + if (need_redis_hash.empty()) { + return ret; + } + + // Second pass: Redis global index (batched). + // We use HMGET to avoid per-hash RTT. + if (redis_lock_ && redis_lock_->connect()) { + auto vals = redis_lock_->hmget(redis_lock_->globalIndexKey(), need_redis_hash); + if (vals.has_value() && vals->size() == need_redis_hash.size()) { + for (size_t j = 0; j < vals->size(); j++) { + const std::string &s = (*vals)[j]; + if (s.empty()) { + continue; + } + const auto pos = s.find(':'); + if (pos == std::string::npos) { + continue; + } + try { + const size_t shard_id = static_cast(std::stoull(s.substr(0, pos))); + if (shard_id >= shard_) { + continue; + } + ret[need_redis_idx[j]] = true; + } catch (...) { + continue; + } + } + return ret; + } + // If Redis IO/parsing fails, fall through to local scan. + } + + // Last resort: scan local in-memory indices. + for (size_t j = 0; j < need_redis_hash.size(); j++) { + const std::string &hash = need_redis_hash[j]; + for (size_t sid = 0; sid < shard_; sid++) { + if (caches_[sid]->exists(hash)) { + ret[need_redis_idx[j]] = true; + break; + } + } + } + + return ret; +} + +size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { + const bool redis_connected = (redis_lock_ && redis_lock_->connect()); + const bool do_global_dedupe = online_mode_ && redis_connected; + const std::string global_key = do_global_dedupe ? redis_lock_->globalIndexKey() : std::string(); + const std::string lock_key = do_global_dedupe ? redis_lock_->hashLockKey(hash) : std::string(); + + // In distributed mode, enable Redis-backed global dedupe + per-hash lock. + // In default single-node mode, avoid extra Redis RTTs on the write hot path. + if (do_global_dedupe) { + // Fast-path: try a single EVAL to reduce RTT. + // NOTE: This may fail on Redis Cluster if keys hash to different slots. + // In that case, fall back to the original two-command sequence. + static const std::string kLuaCheckAndLock = "if redis.call('HEXISTS', KEYS[1], ARGV[2]) == 1 then return 0 end " + "local ok = redis.call('SET', KEYS[2], '1', 'NX', 'PX', ARGV[1]) " + "if ok then return 1 else return 0 end"; + auto got = redis_lock_->evalInt(kLuaCheckAndLock, {global_key, lock_key}, {"30000", hash}); + if (got.has_value()) { + if (*got != 1) { + return 0; + } + } else { + // Fallback (cluster-compatible): HEXISTS then SET NX PX. + if (redis_lock_->hexists(global_key, hash)) { + return 0; + } + if (!redis_lock_->setStringNxPx(lock_key, "1", 30000)) { + return 0; + } + } + } + + const size_t shard_id = online_mode_ ? pickWritableShard(hash) : getShard(hash); + if (shard_id >= shard_) { + if (do_global_dedupe) { + (void)redis_lock_->del(lock_key); + } + return 0; + } + + uint64_t epoch = 0; + if (!isShardWritable(shard_id, &epoch)) { + if (do_global_dedupe) { + (void)redis_lock_->del(lock_key); + } + return 0; + } + + shard_inflight_[shard_id].fetch_add(1, std::memory_order_relaxed); + + size_t slot_id = 0; + std::string evicted_hash; + + // 1. Acquire slot (LRU) + int result = caches_[shard_id]->acquire_slot(hash, slot_id, evicted_hash); + if (result < 0) { + shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); + if (do_global_dedupe) { + (void)redis_lock_->del(lock_key); + } + return 0; + } + + // Maintain local hint bounds: if we evicted something, drop its hint. + if (!evicted_hash.empty()) { + eraseLocalShardHint(evicted_hash); + } + + // If LRU evicted something, delete its Redis mappings *before* we overwrite the slot. + // This avoids a window where stale Redis mapping points to overwritten data. + if (do_global_dedupe && !evicted_hash.empty()) { + // NOTE: best-effort; a Redis failure here can lead to stale mapping. + (void)redis_lock_->hdel(redis_lock_->shardIndexKey(shard_id), evicted_hash); + (void)redis_lock_->hdel(global_key, evicted_hash); + (void)redis_lock_->hdel(redis_lock_->globalCrcKey(), evicted_hash); + } + + const size_t offset_bytes = slot_id * block_size_; + const uint64_t write_offset = static_cast(offset_bytes); + const uint32_t write_len = static_cast(block_size_); + uint32_t data_crc = 0; + + // 2. Write Data to data file (Overwrite) + try { + std::unique_lock lock(*io_locks_[shard_id]); + + if (!pwriteAll(file_fds_[shard_id], buf, block_size_, static_cast(offset_bytes))) { + throw std::runtime_error("Failed to write data"); + } + if (cache::utils::fdatasync_compat(file_fds_[shard_id]) != 0) { + throw std::runtime_error("Failed to fdatasync data"); + } + + // Compute CRC after data is durable. + data_crc = ::crc32(0, reinterpret_cast(buf), block_size_); + + } catch (...) { + caches_[shard_id]->remove(hash); + eraseLocalShardHint(hash); + shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); + if (do_global_dedupe) { + (void)redis_lock_->del(lock_key); + } + return 0; + } + + // Fence again after data is durable (handles revocation mid-write). + uint64_t epoch2 = 0; + if (!isShardWritable(shard_id, &epoch2) || epoch2 != epoch) { + caches_[shard_id]->remove(hash); + eraseLocalShardHint(hash); + shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); + if (do_global_dedupe) { + (void)redis_lock_->del(lock_key); + } + return 0; + } + + // 3. Enqueue journal+redis update; handled by the per-shard journal worker. + auto task = std::make_shared(); + task->shard_id = shard_id; + task->epoch = epoch; + task->write_offset = write_offset; + task->write_len = write_len; + task->data_crc = data_crc; + task->slot_id = slot_id; + task->hash = hash; + task->evicted_hash = evicted_hash; + + { + std::lock_guard lk(*journal_mu_[shard_id]); + journal_queue_[shard_id].push_back(task); + } + journal_cv_[shard_id]->notify_one(); + + // Synchronous commit: wait until journal worker makes the record durable. + { + std::unique_lock lk(task->mu); + task->cv.wait(lk, [&] { return task->done; }); + } + + if (!task->success) { + caches_[shard_id]->remove(hash); + eraseLocalShardHint(hash); + shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); + if (do_global_dedupe) { + (void)redis_lock_->del(lock_key); + } + return 0; + } + + // Mark as ready only after the WAL commit finishes. + caches_[shard_id]->mark_ready(hash, data_crc); + noteLocalShardHint(hash, shard_id); + + shard_written_bytes_[shard_id].fetch_add(static_cast(block_size_), std::memory_order_relaxed); + + shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); + if (do_global_dedupe) { + (void)redis_lock_->del(lock_key); + } + +#ifndef __APPLE__ + if (file_fds_[shard_id] >= 0) { + posix_fadvise(file_fds_[shard_id], offset_bytes, block_size_, POSIX_FADV_DONTNEED); + } +#endif + + return block_size_; +} + +size_t LocalStorageEngine::read(char *buf, const std::string &hash) { + size_t shard_id = static_cast(-1); + size_t slot_id = static_cast(-1); + bool local_hit = false; + bool redis_hit = false; + bool crc_available = false; + uint32_t expected_crc = 0; + + if (!online_mode_) { + shard_id = getShard(hash); + if (!caches_[shard_id]->get_offset_and_crc(hash, slot_id, expected_crc)) { + return 0; + } + crc_available = (expected_crc != 0); + } else { + // 1) Local-first: try hint then (if needed) scan local indices. + auto hinted = localShardHint(hash); + if (hinted.has_value() && hinted.value() < shard_) { + const size_t i = hinted.value(); + const size_t off = caches_[i]->get_offset(hash); + if (off != static_cast(-1)) { + shard_id = i; + slot_id = off; + local_hit = true; + } else { + eraseLocalShardHint(hash); + } + } + + if (shard_id == static_cast(-1)) { + for (size_t i = 0; i < shard_; i++) { + const size_t off = caches_[i]->get_offset(hash); + if (off != static_cast(-1)) { + shard_id = i; + slot_id = off; + local_hit = true; + noteLocalShardHint(hash, i); + break; + } + } + } + + // 2) Fallback: resolve via Redis global index. + if (shard_id == static_cast(-1) && redis_lock_ && redis_lock_->connect()) { + auto resolved = findShardInRedis(hash, &slot_id); + if (resolved.has_value()) { + shard_id = *resolved; + noteLocalShardHint(hash, shard_id); + redis_hit = true; + } + } + } + + if (shard_id == static_cast(-1) || slot_id == static_cast(-1)) { + return 0; + } + + std::shared_lock lock(*io_locks_[shard_id]); + + // Owner-local fast path (online/distributed mode): + // If the mapping is from our local in-memory index AND we currently own the shard, + // then no other node is allowed to overwrite slots in this shard. + // With the shard io lock held, re-check the local offset to avoid races with local eviction. + if (online_mode_ && local_hit) { + const size_t slot_local = caches_[shard_id]->get_offset(hash); + if (slot_local != slot_id) { + eraseLocalShardHint(hash); + return 0; + } + noteLocalShardHint(hash, shard_id); + const size_t offset_bytes = slot_id * block_size_; + if (file_fds_[shard_id] < 0) { + return 0; + } + if (!preadAll(file_fds_[shard_id], buf, block_size_, static_cast(offset_bytes))) { + std::fprintf(stderr, "[light_mem error] read: I/O error for hash %s at offset %zu\n", hash.c_str(), + offset_bytes); + return 0; + } + return block_size_; + } + + // If mapping came from Redis (or Redis is available), re-check it after taking the shard lock. + // This prevents returning wrong data if the owner evicted/overwrote the slot between lookup and read. + if (redis_lock_ && redis_lock_->connect()) { + size_t slot2 = static_cast(-1); + auto shard2 = findShardInRedis(hash, &slot2); + if (!shard2.has_value() || *shard2 != shard_id || slot2 != slot_id) { + return 0; + } + } + + // If Redis was used to resolve, require CRC to be present before reading data. + if (redis_hit && redis_lock_ && redis_lock_->connect()) { + auto crc_s = redis_lock_->hget(redis_lock_->globalCrcKey(), hash); + if (!crc_s.has_value() || crc_s->empty()) { + (void)redis_lock_->hdel(redis_lock_->globalIndexKey(), hash); + (void)redis_lock_->hdel(redis_lock_->shardIndexKey(shard_id), hash); + (void)redis_lock_->hdel(redis_lock_->globalCrcKey(), hash); + return 0; + } + try { + expected_crc = static_cast(std::stoul(*crc_s)); + crc_available = true; + } catch (...) { + (void)redis_lock_->hdel(redis_lock_->globalIndexKey(), hash); + (void)redis_lock_->hdel(redis_lock_->shardIndexKey(shard_id), hash); + (void)redis_lock_->hdel(redis_lock_->globalCrcKey(), hash); + return 0; + } + } + + const size_t offset_bytes = slot_id * block_size_; + if (file_fds_[shard_id] < 0) { + return 0; + } + if (!preadAll(file_fds_[shard_id], buf, block_size_, static_cast(offset_bytes))) { + std::fprintf(stderr, "[light_mem error] read: I/O error for hash %s at offset %zu\n", hash.c_str(), offset_bytes); + return 0; + } + + // Local (offline) correctness: guard against a narrow TOCTOU window where a slot can be evicted + // and reused between index lookup and the actual I/O. In offline mode we don't have Redis to + // re-validate mapping, so use the locally recorded CRC (written at commit time) as a cheap + // correctness check. + if (!online_mode_ && crc_available) { + const uint32_t got_crc = compute_crc32(buf, block_size_); + if (got_crc != expected_crc) { + caches_[shard_id]->remove(hash); + eraseLocalShardHint(hash); + return 0; + } + } + + // CRC verification for Redis-resolved reads. + if (redis_hit && crc_available) { + const uint32_t got_crc = compute_crc32(buf, block_size_); + if (got_crc != expected_crc) { + // Stale Redis mapping: delete best-effort and return miss. + (void)redis_lock_->hdel(redis_lock_->globalIndexKey(), hash); + (void)redis_lock_->hdel(redis_lock_->shardIndexKey(shard_id), hash); + (void)redis_lock_->hdel(redis_lock_->globalCrcKey(), hash); + caches_[shard_id]->remove(hash); + eraseLocalShardHint(hash); + return 0; + } + caches_[shard_id]->put_ready(hash, slot_id, expected_crc); + noteLocalShardHint(hash, shard_id); + } + return block_size_; +} + +std::shared_ptr LocalStorageEngine::getHashInfo() { + auto info = std::make_shared(); + info->caches = caches_; + info->io_locks = io_locks_; + return info; +} + +bool LocalStorageEngine::setHashInfo(const std::shared_ptr &info) { + if (!info) { + std::fprintf(stderr, "[light_mem error] setHashInfo: HashInfo is null\n"); + return false; + } + if (info->caches.size() != shard_ || info->io_locks.size() != shard_) { + std::fprintf(stderr, + "[light_mem error] setHashInfo: shard size mismatch (expected %zu, got caches=%zu io_locks=%zu)\n", + shard_, info->caches.size(), info->io_locks.size()); + return false; + } + caches_ = info->caches; + io_locks_ = info->io_locks; + return true; +} + +} // namespace storage +} // namespace cache diff --git a/src/storage/local_storage_engine_recovery.cpp b/src/storage/local_storage_engine_recovery.cpp new file mode 100755 index 0000000..ac8866b --- /dev/null +++ b/src/storage/local_storage_engine_recovery.cpp @@ -0,0 +1,472 @@ +#include "storage/local_storage_engine.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cache { +namespace storage { + +void LocalStorageEngine::recoverShardToRedis(size_t shard_id) { + if (shard_id >= shard_) { + return; + } + RedisClient *redis = redisForShard(shard_id); + if (!redis || !redis->connect()) { + return; + } + + const size_t shard_capacity = (storage_size_ / shard_) / block_size_; + + // Rebuild from local snapshot + WAL. Redis is NOT a source of truth. + // This routine is intended to be called when a node (re-)acquires ownership of a shard. + std::unique_lock lock(*io_locks_[shard_id]); + + // Load mapping from snapshot if present. + std::unordered_map mapping; + std::unordered_map crc_map; + std::unordered_set crc_present; + { + constexpr uint32_t kSnapshotMagic = 0x534E4150; // SNAP + constexpr uint32_t kSnapshotVersion = 1; + + std::stringstream ss; + ss << filename_ << "_" << shard_id << "/index"; + const std::string snap_path = ss.str(); + + int fd = ::open(snap_path.c_str(), O_RDONLY); + if (fd >= 0) { + auto read_all = [&](void *p, size_t n) -> bool { + char *buf = static_cast(p); + size_t left = n; + while (left > 0) { + ssize_t r = ::read(fd, buf, left); + if (r <= 0) { + return false; + } + buf += static_cast(r); + left -= static_cast(r); + } + return true; + }; + + uint32_t magic = 0; + uint32_t version = 0; + uint64_t count = 0; + if (read_all(&magic, sizeof(magic)) && read_all(&version, sizeof(version)) && read_all(&count, sizeof(count)) && + magic == kSnapshotMagic && version == kSnapshotVersion) { + for (uint64_t i = 0; i < count; i++) { + uint32_t hash_len = 0; + if (!read_all(&hash_len, sizeof(hash_len)) || hash_len > 4096) { + break; + } + std::string hash; + hash.resize(hash_len); + if (hash_len > 0 && !read_all(hash.data(), hash_len)) { + break; + } + uint64_t slot = 0; + if (!read_all(&slot, sizeof(slot))) { + break; + } + uint32_t crc = 0; + if (!read_all(&crc, sizeof(crc))) { + break; + } + if (slot < shard_capacity) { + mapping[hash] = static_cast(slot); + crc_map[hash] = crc; + crc_present.insert(hash); + } + } + } + ::close(fd); + } + } + + // Apply WAL tail (after last checkpoint) to mapping. + std::vector ops; + const off_t start_off = + static_cast((superblock_stable_offset_[shard_id] >= META_HEADER_SIZE) ? superblock_stable_offset_[shard_id] + : META_HEADER_SIZE); + scanWalOps(shard_id, shard_capacity, start_off, ops); + for (const auto &op : ops) { + if (!op.evicted.empty()) { + mapping.erase(op.evicted); + crc_map.erase(op.evicted); + crc_present.erase(op.evicted); + } + if (!op.hash.empty() && op.slot_id < shard_capacity) { + mapping[op.hash] = op.slot_id; + crc_map[op.hash] = op.data_crc; + crc_present.insert(op.hash); + } + } + + const std::string key = redis->shardIndexKey(shard_id); + const std::string gkey = redis->globalIndexKey(); + + // Cleanup stale global-index entries for this shard using the previous shard index. + if (auto prev = redis->hgetall(key); prev.has_value() && (prev->size() % 2 == 0)) { + for (size_t i = 0; i + 1 < prev->size(); i += 2) { + const std::string &old_hash = (*prev)[i]; + if (mapping.find(old_hash) == mapping.end()) { + (void)redis->hdel(gkey, old_hash); + (void)redis->hdel(redis->globalCrcKey(), old_hash); + } + } + } + + // Replace shard index with our rebuilt mapping. + (void)redis->del(key); + + std::vector> cmds; + cmds.reserve(mapping.size() * 3 + 1); + + std::vector buf; + buf.resize(block_size_); + + for (const auto &kv : mapping) { + const size_t slot_id = kv.second; + uint32_t crc = 0; + bool crc_ok = false; + if (crc_present.find(kv.first) != crc_present.end()) { + auto it = crc_map.find(kv.first); + if (it != crc_map.end()) { + crc = it->second; + crc_ok = true; + } + } + if (!crc_ok) { + const size_t offset_bytes = slot_id * block_size_; + if (file_fds_[shard_id] >= 0 && + preadAll(file_fds_[shard_id], buf.data(), block_size_, static_cast(offset_bytes))) { + crc = compute_crc32(buf.data(), block_size_); + crc_ok = true; + } + } + + if (!crc_ok) { + // Unable to compute CRC; skip publishing this mapping to avoid missing CRC in Redis. + continue; + } + + cmds.push_back({"HSET", key, kv.first, std::to_string(slot_id)}); + cmds.push_back({"HSET", gkey, kv.first, std::to_string(shard_id) + ":" + std::to_string(slot_id)}); + cmds.push_back({"HSET", redis->globalCrcKey(), kv.first, std::to_string(crc)}); + } + cmds.push_back({"SET", redis->shardSeqKey(shard_id), std::to_string(superblock_seq_[shard_id])}); + if (!redis->pipeline(cmds)) { + throw std::runtime_error("Recovery Redis pipeline failed"); + } +} + +bool LocalStorageEngine::shouldFullRecoverRedis(size_t shard_id) { + if (shard_id >= shard_) { + return false; + } + RedisClient *redis = redisForShard(shard_id); + if (!redis || !redis->connect()) { + return false; + } + + const std::string key = redis->shardIndexKey(shard_id); + + auto hlen = redis->hlen(key); + if (!hlen.has_value() || *hlen == 0) { + return true; + } + + auto rseq = redis->getString(redis->shardSeqKey(shard_id)); + if (!rseq.has_value() || rseq->empty()) { + return true; + } + try { + const uint64_t seq = static_cast(std::stoull(*rseq)); + if (seq != superblock_seq_[shard_id]) { + return true; + } + } catch (...) { + return true; + } + + return false; +} + +void LocalStorageEngine::recoverShardToRedisIncremental(size_t shard_id) { + if (shard_id >= shard_) { + return; + } + RedisClient *redis = redisForShard(shard_id); + if (!redis || !redis->connect()) { + return; + } + + const size_t shard_capacity = (storage_size_ / shard_) / block_size_; + const off_t start_off = + static_cast((superblock_stable_offset_[shard_id] >= META_HEADER_SIZE) ? superblock_stable_offset_[shard_id] + : META_HEADER_SIZE); + + std::vector ops; + scanWalOps(shard_id, shard_capacity, start_off, ops); + if (ops.empty()) { + return; + } + + const std::string key = redis->shardIndexKey(shard_id); + const std::string gkey = redis->globalIndexKey(); + + std::vector> cmds; + cmds.reserve(ops.size() * 5 + 1); + for (const auto &op : ops) { + if (!op.evicted.empty()) { + cmds.push_back({"HDEL", key, op.evicted}); + cmds.push_back({"HDEL", gkey, op.evicted}); + cmds.push_back({"HDEL", redis->globalCrcKey(), op.evicted}); + } + cmds.push_back({"HSET", key, op.hash, std::to_string(op.slot_id)}); + cmds.push_back({"HSET", gkey, op.hash, std::to_string(shard_id) + ":" + std::to_string(op.slot_id)}); + cmds.push_back({"HSET", redis->globalCrcKey(), op.hash, std::to_string(op.data_crc)}); + } + cmds.push_back({"SET", redis->shardSeqKey(shard_id), std::to_string(superblock_seq_[shard_id])}); + + (void)redis->pipeline(cmds); +} + +void LocalStorageEngine::recoverShardToRedisSmart(size_t shard_id) { + if (shouldFullRecoverRedis(shard_id)) { + recoverShardToRedis(shard_id); + return; + } + recoverShardToRedisIncremental(shard_id); +} + +void LocalStorageEngine::scanWalOps(size_t shard_id, size_t shard_capacity, off_t start_off, + std::vector &ops) { + struct stat st{}; + if (::fstat(meta_fds_[shard_id], &st) != 0) { + return; + } + const off_t end = st.st_size; + off_t off = start_off; + + uint64_t max_epoch_seen = 0; + + while (off + static_cast(sizeof(uint32_t)) <= end) { + uint32_t magic = 0; + if (!preadAll(meta_fds_[shard_id], &magic, sizeof(uint32_t), off)) { + break; + } + if (magic != JOURNAL_MAGIC) { + break; + } + + // Parse journal record (epoch-aware). + bool parsed = false; + if (off + static_cast(sizeof(JournalRecord)) <= end) { + JournalRecord rec{}; + if (preadAll(meta_fds_[shard_id], &rec, sizeof(JournalRecord), off) && rec.magic == JOURNAL_MAGIC && + rec.version == JOURNAL_VERSION && (rec.flags & ~JOURNAL_FLAG_EPOCH_MARKER) == 0 && rec.hash_len <= 4096 && + rec.evicted_hash_len <= 4096) { + off += static_cast(sizeof(JournalRecord)); + if (off + static_cast(rec.hash_len) + static_cast(rec.evicted_hash_len) + + static_cast(sizeof(uint32_t)) > + end) { + break; + } + + std::string hash; + hash.resize(rec.hash_len); + if (rec.hash_len > 0) { + if (!preadAll(meta_fds_[shard_id], hash.data(), rec.hash_len, off)) { + break; + } + } + off += static_cast(rec.hash_len); + + std::string evicted; + evicted.resize(rec.evicted_hash_len); + if (rec.evicted_hash_len > 0) { + if (!preadAll(meta_fds_[shard_id], evicted.data(), rec.evicted_hash_len, off)) { + break; + } + } + off += static_cast(rec.evicted_hash_len); + + uint32_t record_crc = 0; + if (!preadAll(meta_fds_[shard_id], &record_crc, sizeof(uint32_t), off)) { + break; + } + off += static_cast(sizeof(uint32_t)); + + std::vector tmp; + tmp.resize(sizeof(JournalRecord) + hash.size() + evicted.size()); + std::memcpy(tmp.data(), &rec, sizeof(JournalRecord)); + if (!hash.empty()) { + std::memcpy(tmp.data() + sizeof(JournalRecord), hash.data(), hash.size()); + } + if (!evicted.empty()) { + std::memcpy(tmp.data() + sizeof(JournalRecord) + hash.size(), evicted.data(), evicted.size()); + } + const uint32_t expect = compute_crc32(tmp.data(), tmp.size()); + if (expect != record_crc) { + parsed = true; + continue; + } + + if ((rec.flags & JOURNAL_FLAG_EPOCH_MARKER) != 0) { + if (rec.epoch >= max_epoch_seen) { + max_epoch_seen = rec.epoch; + } + parsed = true; + continue; + } + + if (rec.epoch < max_epoch_seen) { + parsed = true; + continue; + } + if (rec.epoch > max_epoch_seen) { + max_epoch_seen = rec.epoch; + } + + if (rec.write_offset % block_size_ != 0) { + parsed = true; + continue; + } + const size_t slot_id = static_cast(rec.write_offset / block_size_); + if (slot_id >= shard_capacity) { + parsed = true; + continue; + } + + JournalOp op; + op.hash = std::move(hash); + op.evicted = std::move(evicted); + op.slot_id = slot_id; + op.data_crc = rec.data_crc; + ops.emplace_back(std::move(op)); + parsed = true; + } + } + + if (parsed) { + continue; + } + + // Not a valid record at this offset; stop scanning. + break; + } +} + +void LocalStorageEngine::recoverAllShards(size_t shard_capacity) { + for (size_t i = 0; i < shard_; i++) { + recoverShard(i, shard_capacity); + } +} + +void LocalStorageEngine::recoverShard(size_t shard_id, size_t shard_capacity) { + RedisClient *redis = redisForShard(shard_id); + const bool redis_ok = (redis && redis->connect()); + caches_[shard_id]->reset(); + + // 0. Try to load from Local Snapshot first + std::stringstream ss; + ss << filename_ << "_" << shard_id << "/index"; + std::string snap_path = ss.str(); + bool snapshot_loaded = caches_[shard_id]->loadFromSnapshot(snap_path); + + // If snapshot is loaded, also warm up the local hash->shard hint map so distributed-mode + // reads/queries can locate the shard in O(1) without scanning all shards or hitting Redis. + if (snapshot_loaded && online_mode_) { + std::vector hashes; + caches_[shard_id]->dump_ready(hashes); + for (const auto &hash : hashes) { + noteLocalShardHint(hash, shard_id); + } + } + + // 1. Scan WAL + // We rely SOLELY on Snapshot + WAL. + // Redis is treated as a cache/index that reflects our state, not the source of truth. + const off_t start_off = + static_cast((superblock_stable_offset_[shard_id] >= META_HEADER_SIZE) ? superblock_stable_offset_[shard_id] + : META_HEADER_SIZE); + + std::vector ops; + scanWalOps(shard_id, shard_capacity, start_off, ops); + + // 2. Apply WAL to Memory + if (!ops.empty()) { + for (const auto &op : ops) { + if (!op.evicted.empty()) { + caches_[shard_id]->remove(op.evicted); + if (online_mode_) { + eraseLocalShardHint(op.evicted); + } + } + caches_[shard_id]->put_ready(op.hash, op.slot_id, op.data_crc); + if (online_mode_) { + noteLocalShardHint(op.hash, shard_id); + } + } + } + + // 3. Sync WAL updates to Redis (and truncate WAL if successful) + if (!ops.empty() && redis_ok) { + writeCheckpointSuperblockOnly(shard_id); + + bool replay_ok = true; + const std::string key = redis->shardIndexKey(shard_id); + const std::string gkey = redis->globalIndexKey(); + for (const auto &op : ops) { + if (!op.evicted.empty()) { + if (!redis->hdel(key, op.evicted)) { + replay_ok = false; + break; + } + (void)redis->hdel(gkey, op.evicted); + (void)redis->hdel(redis->globalCrcKey(), op.evicted); + } + if (!redis->hset(key, op.hash, std::to_string(op.slot_id))) { + replay_ok = false; + break; + } + + if (!redis->hset(redis->globalCrcKey(), op.hash, std::to_string(op.data_crc))) { + replay_ok = false; + break; + } + + // Global mapping: only recover if missing; duplicates are intentionally ignored. + (void)redis->hsetnx(gkey, op.hash, std::to_string(shard_id) + ":" + std::to_string(op.slot_id)); + } + if (replay_ok) { + if (!redis->setString(redis->shardSeqKey(shard_id), std::to_string(superblock_seq_[shard_id]))) { + replay_ok = false; + } + } + + if (replay_ok) { + truncateJournalToHeader(shard_id); + } + } + + // 4. Handle Empty State / Fresh Shard + // If we have no snapshot and no WAL logs, we are effectively empty. + // We must ensure Redis doesn't hold stale data pointing to us. + if (!snapshot_loaded && ops.empty() && redis_ok) { + (void)redis->del(redis->shardIndexKey(shard_id)); + } +} + +} // namespace storage +} // namespace cache diff --git a/src/storage/local_storage_engine_shard.cpp b/src/storage/local_storage_engine_shard.cpp new file mode 100755 index 0000000..103aa38 --- /dev/null +++ b/src/storage/local_storage_engine_shard.cpp @@ -0,0 +1,236 @@ +#include "storage/local_storage_engine.h" + +#include +#include + +namespace cache { +namespace storage { + +namespace { + +// Small, fast 64-bit mixer (SplitMix64). Suitable for generating a pseudo-random +// probe sequence from a stable seed. +static inline uint64_t splitmix64(uint64_t x) { + x += 0x9e3779b97f4a7c15ull; + x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9ull; + x = (x ^ (x >> 27)) * 0x94d049bb133111ebull; + return x ^ (x >> 31); +} + +} // namespace + +size_t LocalStorageEngine::getShard(const std::string &hash) const { return std::hash{}(hash) % shard_; } + +bool LocalStorageEngine::isShardWritable(size_t shard_id, uint64_t *epoch_out) const { + if (shard_id >= shard_) { + return false; + } + if (shard_writable_[shard_id].load(std::memory_order_relaxed) == 0) { + return false; + } + if (shard_draining_[shard_id].load(std::memory_order_relaxed) != 0) { + return false; + } + if (epoch_out) { + *epoch_out = shard_epoch_cache_[shard_id].load(std::memory_order_relaxed); + } + return true; +} + +std::optional LocalStorageEngine::findShardInRedis(const std::string &hash, size_t *slot_id_out) { + if (!redis_lock_ || !redis_lock_->connect()) { + return std::nullopt; + } + const auto v = redis_lock_->hget(redis_lock_->globalIndexKey(), hash); + if (!v.has_value()) { + return std::nullopt; + } + const auto &s = *v; + const auto pos = s.find(':'); + if (pos == std::string::npos) { + return std::nullopt; + } + try { + const size_t shard_id = static_cast(std::stoull(s.substr(0, pos))); + const size_t slot_id = static_cast(std::stoull(s.substr(pos + 1))); + if (shard_id >= shard_) { + return std::nullopt; + } + if (slot_id_out) { + *slot_id_out = slot_id; + } + return shard_id; + } catch (...) { + return std::nullopt; + } +} + +size_t LocalStorageEngine::pickWritableShard(const std::string &hash) const { + if (shard_ == 0) { + return 0; + } + + // Prefer a stable mapping (hash -> shard). If that shard is not currently writable + // (ownership/draining), DO NOT fall back in ring order: that can amplify load on the + // next writable shard after long non-writable runs, causing severe write skew and + // premature per-shard eviction when ownership is split across nodes. + // + // Instead, probe pseudo-random shard IDs derived from the hash until we find a writable shard. + // This keeps writes approximately uniform over the writable set. + const size_t preferred = getShard(hash); + if (isShardWritable(preferred, nullptr)) { + return preferred; + } + + // Seed from std::hash (process-stable). We only need per-process stability here. + uint64_t seed = static_cast(std::hash{}(hash)); + seed = splitmix64(seed ^ 0x6a09e667f3bcc909ull); + + // Fast-path: try a small number of randomized probes (expected to succeed quickly when a + // non-trivial fraction of shards are writable). + const size_t max_probes = (shard_ < 16) ? shard_ : 16; + for (size_t i = 0; i < max_probes; i++) { + seed = splitmix64(seed); + const size_t sid = static_cast(seed % static_cast(shard_)); + if (isShardWritable(sid, nullptr)) { + return sid; + } + } + + // Slow-path: guarantee progress. + for (size_t sid = 0; sid < shard_; sid++) { + if (isShardWritable(sid, nullptr)) { + return sid; + } + } + return shard_; +} + +void LocalStorageEngine::updateShardAssignments(const std::vector &shard_ids, + const std::vector &epochs, + const std::vector &draining) { + if (shard_ids.size() != epochs.size() || shard_ids.size() != draining.size()) { + std::fprintf(stderr, "[light_mem error] updateShardAssignments: size mismatch (ids=%zu epochs=%zu draining=%zu)\n", + shard_ids.size(), epochs.size(), draining.size()); + return; + } + + // Build desired state. + std::vector desired_writable(shard_, 0); + std::vector desired_draining(shard_, 1); // default: treat as draining (no new writes) + std::vector desired_epoch(shard_, 0); + + for (size_t i = 0; i < shard_ids.size(); i++) { + const size_t sid = shard_ids[i]; + if (sid >= shard_) { + continue; + } + desired_writable[sid] = 1; + desired_draining[sid] = (draining[i] != 0) ? 1 : 0; + desired_epoch[sid] = epochs[i]; + } + + const size_t shard_capacity = (storage_size_ / shard_) / block_size_; + + for (size_t sid = 0; sid < shard_; sid++) { + const uint8_t new_w = desired_writable[sid]; + const uint8_t new_d = desired_draining[sid]; + const uint64_t new_e = desired_epoch[sid]; + + const uint64_t old_e = shard_epoch_cache_[sid].load(std::memory_order_relaxed); + const uint8_t old_w = shard_writable_[sid].load(std::memory_order_relaxed); + + shard_writable_[sid].store(new_w, std::memory_order_relaxed); + shard_draining_[sid].store(new_d, std::memory_order_relaxed); + shard_epoch_cache_[sid].store(new_e, std::memory_order_relaxed); + + // Persist epoch changes for shards that are (now) writable. + if (new_w != 0 && new_e != 0 && new_e != old_e) { + std::unique_lock io_lock(*io_locks_[sid]); + try { + onEpochChangedLocked(sid, new_e); + } catch (const std::exception &e) { + std::fprintf(stderr, "[light_mem warning] updateShardAssignments: failed to persist epoch for shard %zu: %s\n", + sid, e.what()); + } catch (...) { + std::fprintf( + stderr, + "[light_mem warning] updateShardAssignments: failed to persist epoch for shard %zu: unknown error\n", sid); + } + } + + // If shard ownership is lost, clear local index and related hints. + if (old_w != 0 && new_w == 0) { + if (caches_[sid]) { + caches_[sid]->reset(); + } + shard_recovered_epoch_[sid] = 0; + for (size_t b = 0; b < kLocalHintBuckets; b++) { + std::unique_lock lk(local_hint_mu_[b]); + auto &bucket = local_hash_to_shard_[b]; + for (auto it = bucket.begin(); it != bucket.end();) { + if (it->second == sid) { + it = bucket.erase(it); + } else { + ++it; + } + } + } + } + + // In online mode, rebuild local index only for owned/writable shards when ownership changes. + if (online_mode_ && new_w != 0 && new_e != 0) { + const bool need_recover = (old_w == 0) || (new_e != old_e) || (shard_recovered_epoch_[sid] != new_e); + if (need_recover) { + try { + recoverShard(sid, shard_capacity); + shard_recovered_epoch_[sid] = new_e; + } catch (...) { + // Best-effort: leave shard_recovered_epoch_ unchanged on failures. + } + } + } + } +} + +uint32_t LocalStorageEngine::shardInflight(size_t shard_id) const { + if (shard_id >= shard_) { + return 0; + } + return shard_inflight_[shard_id].load(std::memory_order_relaxed); +} + +uint64_t LocalStorageEngine::shardWrittenBytes(size_t shard_id) const { + if (shard_id >= shard_) { + return 0; + } + return shard_written_bytes_[shard_id].load(std::memory_order_relaxed); +} + +uint64_t LocalStorageEngine::shardEvictionCount(size_t shard_id) const { + if (shard_id >= shard_) { + return 0; + } + const auto &c = caches_[shard_id]; + if (!c) { + return 0; + } + return c->eviction_count(); +} + +uint64_t LocalStorageEngine::evictionCount() const { + uint64_t total = 0; + for (size_t i = 0; i < shard_; i++) { + const auto &c = caches_[i]; + if (!c) { + continue; + } + total += c->eviction_count(); + } + return total; +} + +bool LocalStorageEngine::evictionObserved() const { return evictionCount() > 0; } + +} // namespace storage +} // namespace cache diff --git a/src/storage/local_storage_wal.h b/src/storage/local_storage_wal.h new file mode 100755 index 0000000..e35f202 --- /dev/null +++ b/src/storage/local_storage_wal.h @@ -0,0 +1,62 @@ +#pragma once + +#include +#include + +#include // crc32 + +namespace cache { +namespace storage { + +#pragma pack(push, 1) + +struct SuperBlock { + uint32_t magic; // 0x4C4D454D (LMEM) + uint32_t version; // Format version + uint64_t shard_id; // Shard ID + uint64_t sequence_id; // Monotonically increasing ID + uint64_t stable_offset; // Offset in meta file where Journal scanning starts (truncate-mode WAL) + uint64_t total_capacity; // Max capacity of the shard + uint32_t block_size; // Block size + uint64_t current_epoch; // Epoch for fencing + uint8_t padding[4096 - 56]; + uint32_t crc32; // CRC of the first 4092 bytes +}; + +// Journal record: +// - Includes epoch for fencing. +// - Supports an epoch-marker record (flags & kFlagEpochMarker) to establish epoch boundaries in the append stream. +// Variable bytes layout: [hash bytes][evicted_hash bytes][record_crc32] +struct JournalRecord { + uint32_t magic; // 0x4A524E4C (JRNL) + uint16_t version; // Format version + uint16_t flags; // bitmask + uint64_t epoch; // Fencing epoch (monotonic per shard ownership) + uint64_t write_offset; // Offset in data file + uint32_t write_len; // Length of data written + uint32_t data_crc; // CRC of the data + uint32_t hash_len; // Length of hash key + uint32_t evicted_hash_len; // Length of evicted hash (0 if none) +#if defined(__GNUC__) || defined(__clang__) + uint8_t hash_key[0]; // Variable-length bytes follow this header (compiler extension) +#endif +}; + +#pragma pack(pop) + +inline constexpr uint32_t SUPERBLOCK_MAGIC = 0x4C4D454D; +inline constexpr uint32_t SUPERBLOCK_VERSION = 1; +inline constexpr uint32_t JOURNAL_MAGIC = 0x4A524E4C; +inline constexpr uint16_t JOURNAL_VERSION = 1; +inline constexpr uint16_t JOURNAL_FLAG_EPOCH_MARKER = 1u << 0; +inline constexpr size_t SUPERBLOCK_SIZE = 4096; +inline constexpr size_t META_HEADER_SIZE = 8192; // 2 * SuperBlock + +static_assert(sizeof(SuperBlock) == SUPERBLOCK_SIZE, "SuperBlock must be exactly 4096 bytes"); + +inline uint32_t compute_crc32(const void *data, size_t len) { + return ::crc32(0, reinterpret_cast(data), len); +} + +} // namespace storage +} // namespace cache diff --git a/src/storage/redis_client.h b/src/storage/redis_client.h new file mode 100755 index 0000000..6f8f648 --- /dev/null +++ b/src/storage/redis_client.h @@ -0,0 +1,545 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cache::storage { + +class RedisClient { +public: + struct Options { + std::string host = "127.0.0.1"; + int port = 6379; + std::string password; + int db = 0; + std::string key_prefix; + int connect_timeout_ms = 300; + // NOTE: Redis metadata updates can be pipelined in large batches. + // Keep this comfortably above typical batch drain time to avoid + // silent timeouts that desynchronize the connection. + int io_timeout_ms = 5000; + }; + + explicit RedisClient(Options opt) : opt_(std::move(opt)) {} + + ~RedisClient() { close(); } + + RedisClient(const RedisClient &) = delete; + RedisClient &operator=(const RedisClient &) = delete; + + bool connect() { + std::lock_guard lk(mu_); + if (sock_ >= 0) { + return true; + } + + addrinfo hints{}; + hints.ai_family = AF_UNSPEC; + hints.ai_socktype = SOCK_STREAM; + + addrinfo *res = nullptr; + const std::string port_str = std::to_string(opt_.port); + if (::getaddrinfo(opt_.host.c_str(), port_str.c_str(), &hints, &res) != 0) { + return false; + } + + int sock = -1; + for (addrinfo *p = res; p != nullptr; p = p->ai_next) { + sock = ::socket(p->ai_family, p->ai_socktype, p->ai_protocol); + if (sock < 0) { + continue; + } + + setSocketOptions(sock); + + if (::connect(sock, p->ai_addr, p->ai_addrlen) == 0) { + break; + } + + ::close(sock); + sock = -1; + } + ::freeaddrinfo(res); + + if (sock < 0) { + return false; + } + + sock_ = sock; + + // AUTH + if (!opt_.password.empty()) { + auto resp = command({"AUTH", opt_.password}); + if (!resp.ok) { + close(); + return false; + } + } + + // SELECT DB + if (opt_.db != 0) { + auto resp = command({"SELECT", std::to_string(opt_.db)}); + if (!resp.ok) { + close(); + return false; + } + } + + // PING + auto ping = command({"PING"}); + if (!ping.ok) { + close(); + return false; + } + + return true; + } + + void close() { + std::lock_guard lk(mu_); + if (sock_ >= 0) { + ::close(sock_); + sock_ = -1; + } + } + + bool isConnected() const { return sock_ >= 0; } + + bool setString(const std::string &key, const std::string &value) { + auto resp = command({"SET", key, value}); + return resp.ok; + } + + std::optional getString(const std::string &key) { + auto resp = command({"GET", key}); + if (!resp.ok) { + return std::nullopt; + } + return resp.bulk; + } + + bool hset(const std::string &key, const std::string &field, const std::string &value) { + auto resp = command({"HSET", key, field, value}); + return resp.ok; + } + + bool hsetnx(const std::string &key, const std::string &field, const std::string &value) { + auto resp = command({"HSETNX", key, field, value}); + if (!resp.ok || !resp.bulk.has_value()) { + return false; + } + // Integer reply: 1 if field is a new field, 0 if it was already present. + try { + return std::stoll(*resp.bulk) == 1; + } catch (...) { + return false; + } + } + + bool hexists(const std::string &key, const std::string &field) { + auto resp = command({"HEXISTS", key, field}); + if (!resp.ok || !resp.bulk.has_value()) { + return false; + } + try { + return std::stoll(*resp.bulk) == 1; + } catch (...) { + return false; + } + } + + std::optional hget(const std::string &key, const std::string &field) { + auto resp = command({"HGET", key, field}); + if (!resp.ok) { + return std::nullopt; + } + return resp.bulk; + } + + std::optional hlen(const std::string &key) { + auto resp = command({"HLEN", key}); + if (!resp.ok || !resp.bulk.has_value()) { + return std::nullopt; + } + try { + return static_cast(std::stoull(*resp.bulk)); + } catch (...) { + return std::nullopt; + } + } + + // HMGET key field1 field2 ... + // Returns one entry per requested field. Missing fields are returned as empty strings. + // NOTE: this matches parseResp() behavior for nil bulk strings inside arrays. + std::optional> hmget(const std::string &key, const std::vector &fields) { + if (fields.empty()) { + return std::vector{}; + } + std::vector argv; + argv.reserve(2 + fields.size()); + argv.emplace_back("HMGET"); + argv.emplace_back(key); + for (const auto &f : fields) { + argv.emplace_back(f); + } + auto resp = command(argv); + if (!resp.ok) { + return std::nullopt; + } + // RESP array length should match requested fields. + if (resp.array.size() != fields.size()) { + return std::nullopt; + } + return resp.array; + } + + bool hdel(const std::string &key, const std::string &field) { + auto resp = command({"HDEL", key, field}); + return resp.ok; + } + + bool del(const std::string &key) { + auto resp = command({"DEL", key}); + return resp.ok; + } + + // Send multiple commands back-to-back (RESP pipelining) and parse all replies. + // Returns true if the I/O succeeded; individual Redis error replies are tolerated. + // This is intended for high-throughput metadata updates (e.g., HSET/HDEL) where + // callers historically ignored per-command success. + bool pipeline(const std::vector> &argv_list) { + std::lock_guard lk(mu_); + if (sock_ < 0) { + return false; + } + if (argv_list.empty()) { + return true; + } + + std::string req; + req.reserve(argv_list.size() * 64); + for (const auto &argv : argv_list) { + req += "*" + std::to_string(argv.size()) + "\r\n"; + for (const auto &a : argv) { + req += "$" + std::to_string(a.size()) + "\r\n"; + req += a; + req += "\r\n"; + } + } + + if (!writeAll(req.data(), req.size())) { + close(); + return false; + } + + // Drain replies to keep the connection in sync. + for (size_t i = 0; i < argv_list.size(); i++) { + (void)parseResp(); + if (sock_ < 0) { + return false; + } + } + return true; + } + + // EVAL helper (used to reduce round-trips on hot paths). + std::optional evalInt(const std::string &script, const std::vector &keys, + const std::vector &args) { + std::vector argv; + argv.reserve(3 + keys.size() + args.size()); + argv.emplace_back("EVAL"); + argv.emplace_back(script); + argv.emplace_back(std::to_string(keys.size())); + for (const auto &k : keys) { + argv.emplace_back(k); + } + for (const auto &a : args) { + argv.emplace_back(a); + } + auto resp = command(argv); + if (!resp.ok || !resp.bulk.has_value()) { + return std::nullopt; + } + try { + return std::stoll(*resp.bulk); + } catch (...) { + return std::nullopt; + } + } + + // SET key value NX PX ttl_ms + bool setStringNxPx(const std::string &key, const std::string &value, uint64_t ttl_ms) { + auto resp = command({"SET", key, value, "NX", "PX", std::to_string(ttl_ms)}); + // Success returns +OK; failure returns nil bulk. + return resp.ok && resp.bulk.has_value(); + } + + // Returns alternating field/value list + std::optional> hgetall(const std::string &key) { + auto resp = command({"HGETALL", key}); + if (!resp.ok) { + return std::nullopt; + } + return resp.array; + } + + std::string shardIndexKey(size_t shard_id) const { + return opt_.key_prefix + ":" + std::to_string(shard_id) + ":index"; + } + + // Global hash -> "shard_id:slot_id" mapping. + std::string globalIndexKey() const { return opt_.key_prefix + ":global:index"; } + + // Global hash -> data CRC (uint32 as string). + std::string globalCrcKey() const { return opt_.key_prefix + ":global:crc"; } + + // Per-hash lock key used to prevent concurrent writers from racing. + std::string hashLockKey(const std::string &hash) const { return opt_.key_prefix + ":lock:" + hash; } + + std::string shardSeqKey(size_t shard_id) const { return opt_.key_prefix + ":" + std::to_string(shard_id) + ":seq"; } + +private: + struct Resp { + bool ok = false; + std::optional bulk; + std::vector array; + }; + + void setSocketOptions(int sock) { + int yes = 1; + ::setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &yes, sizeof(yes)); + + timeval tv{}; + tv.tv_sec = opt_.io_timeout_ms / 1000; + tv.tv_usec = (opt_.io_timeout_ms % 1000) * 1000; + ::setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)); + ::setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)); + } + + bool writeAll(const void *data, size_t len) { + const char *p = static_cast(data); + size_t left = len; + while (left > 0) { + ssize_t n = ::send(sock_, p, left, 0); + if (n <= 0) { + return false; + } + p += static_cast(n); + left -= static_cast(n); + } + return true; + } + + bool readExact(void *buf, size_t len) { + char *p = static_cast(buf); + size_t left = len; + while (left > 0) { + ssize_t n = ::recv(sock_, p, left, 0); + if (n <= 0) { + return false; + } + p += static_cast(n); + left -= static_cast(n); + } + return true; + } + + bool readLine(std::string &out) { + out.clear(); + char c; + while (true) { + if (!readExact(&c, 1)) { + return false; + } + if (c == '\r') { + char lf; + if (!readExact(&lf, 1)) { + return false; + } + if (lf != '\n') { + return false; + } + return true; + } + out.push_back(c); + } + } + + Resp command(const std::vector &argv) { + std::lock_guard lk(mu_); + if (sock_ < 0) { + return {}; + } + + std::string req; + req.reserve(64); + req += "*" + std::to_string(argv.size()) + "\r\n"; + for (const auto &a : argv) { + req += "$" + std::to_string(a.size()) + "\r\n"; + req += a; + req += "\r\n"; + } + + if (!writeAll(req.data(), req.size())) { + close(); + return {}; + } + + return parseResp(); + } + + Resp parseResp() { + char type; + if (!readExact(&type, 1)) { + close(); + return {}; + } + + std::string line; + if (!readLine(line)) { + close(); + return {}; + } + + if (type == '+') { + // Simple string OK + Resp r; + r.ok = true; + r.bulk = line; + return r; + } + + if (type == '-') { + // Error + Resp r; + r.ok = false; + return r; + } + + if (type == ':') { + Resp r; + r.ok = true; + r.bulk = line; + return r; + } + + if (type == '$') { + int64_t n = 0; + try { + n = std::stoll(line); + } catch (...) { + close(); + return {}; + } + if (n < 0) { + Resp r; + r.ok = true; + r.bulk = std::nullopt; + return r; + } + std::string bulk; + bulk.resize(static_cast(n)); + if (!readExact(bulk.data(), bulk.size())) { + close(); + return {}; + } + // consume CRLF + char crlf[2]; + if (!readExact(crlf, 2) || crlf[0] != '\r' || crlf[1] != '\n') { + close(); + return {}; + } + Resp r; + r.ok = true; + r.bulk = std::move(bulk); + return r; + } + + if (type == '*') { + int64_t count = 0; + try { + count = std::stoll(line); + } catch (...) { + close(); + return {}; + } + if (count < 0) { + Resp r; + r.ok = true; + return r; + } + std::vector arr; + arr.reserve(static_cast(count)); + for (int64_t i = 0; i < count; i++) { + char t; + if (!readExact(&t, 1)) { + close(); + return {}; + } + std::string l; + if (!readLine(l)) { + close(); + return {}; + } + if (t == '$') { + int64_t n = 0; + try { + n = std::stoll(l); + } catch (...) { + close(); + return {}; + } + if (n < 0) { + arr.emplace_back(""); + continue; + } + std::string bulk; + bulk.resize(static_cast(n)); + if (!readExact(bulk.data(), bulk.size())) { + close(); + return {}; + } + char crlf[2]; + if (!readExact(crlf, 2) || crlf[0] != '\r' || crlf[1] != '\n') { + close(); + return {}; + } + arr.emplace_back(std::move(bulk)); + } else if (t == ':') { + arr.emplace_back(l); + } else if (t == '+') { + arr.emplace_back(l); + } else { + // Unsupported nested types + close(); + return {}; + } + } + Resp r; + r.ok = true; + r.array = std::move(arr); + return r; + } + + close(); + return {}; + } + + Options opt_; + int sock_ = -1; + mutable std::recursive_mutex mu_; +}; + +} // namespace cache::storage diff --git a/src/storage/storage_engine.h b/src/storage/storage_engine.h old mode 100644 new mode 100755 diff --git a/src/utils/fsync_compat.h b/src/utils/fsync_compat.h new file mode 100755 index 0000000..c9a60ea --- /dev/null +++ b/src/utils/fsync_compat.h @@ -0,0 +1,28 @@ +#pragma once + +#include + +#if defined(__APPLE__) +#include +#endif + +namespace cache { +namespace utils { + +inline int fdatasync_compat(int fd) { +#if defined(__APPLE__) + // macOS doesn't provide fdatasync(2). Best-effort equivalent is F_FULLFSYNC + // (flushes to physical media). If unavailable/fails, fall back to fsync. +#ifdef F_FULLFSYNC + if (::fcntl(fd, F_FULLFSYNC) == 0) { + return 0; + } +#endif + return ::fsync(fd); +#else + return ::fdatasync(fd); +#endif +} + +} // namespace utils +} // namespace cache diff --git a/tests/mixed_lru.py b/tests/mixed_lru.py old mode 100644 new mode 100755 diff --git a/tests/multi_node/harness.py b/tests/multi_node/harness.py new file mode 100755 index 0000000..c40a3cc --- /dev/null +++ b/tests/multi_node/harness.py @@ -0,0 +1,741 @@ +#!/usr/bin/env python3 +"""Multi-node test harness for LightMem. + +Goals: +- Each test file can run standalone: it starts/stops Redis+Etcd as needed. +- run_all.py can start Redis+Etcd once and pass connection info to each test to reuse. + +We keep this harness dependency-light: +- Redis: prefer local redis-server; docker compose fallback. +- Etcd: prefer local etcd; docker (single container) fallback. + +All tests are designed to run on a single machine. +""" + +from __future__ import annotations + +import argparse +import json +import os +import signal +import shutil +import socket +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +import importlib.util +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Optional + + +ROOT = Path(__file__).resolve().parents[2] + +# Ensure repo's Python sources are importable when running tests directly +# without installing the package (editable install still works fine). +_PY_SRC = ROOT / "python" +if _PY_SRC.exists(): + sys.path.insert(0, str(_PY_SRC)) + + +def _load_etcd_http_client_cls(): + """Load EtcdV3HttpClient from repo sources without requiring installed package.""" + path = ROOT / "python" / "light_mem" / "etcd_v3_http.py" + if not path.exists(): + return None + spec = importlib.util.spec_from_file_location("light_mem_etcd_v3_http", str(path)) + if spec is None or spec.loader is None: + return None + module = importlib.util.module_from_spec(spec) + # Register into sys.modules before exec_module so decorators (e.g. dataclasses) + # can resolve module globals via sys.modules[__module__]. + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return getattr(module, "EtcdV3HttpClient", None) + + +def which(name: str) -> str | None: + return shutil.which(name) + + +def run(cmd: list[str], *, cwd: Path | None = None, env: dict[str, str] | None = None, timeout: int = 180) -> None: + subprocess.run(cmd, cwd=str(cwd) if cwd else None, env=env, check=True, timeout=timeout) + + +def find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return int(s.getsockname()[1]) + + +def wait_port(host: str, port: int, *, timeout_s: float = 15.0) -> bool: + deadline = time.time() + timeout_s + while time.time() < deadline: + try: + with socket.create_connection((host, port), timeout=0.2): + return True + except OSError: + time.sleep(0.1) + return False + + +def redis_ping(host: str, port: int, *, timeout_s: float = 0.5) -> bool: + try: + return redis_resp_command(host, port, ["PING"], timeout_s=timeout_s).startswith("+PONG") + except Exception: + return False + + +def etcd_health(host: str, port: int, *, timeout_s: float = 0.8) -> bool: + url = f"http://{host}:{port}/health" + try: + with urllib.request.urlopen(url, timeout=timeout_s) as resp: + raw = resp.read() + obj = json.loads(raw.decode("utf-8")) if raw else {} + return str(obj.get("health", "")).lower() == "true" + except Exception: + return False + + +def _wait_lightmem_server_ready(*, host: str, redis_port: int, etcd_port: int, proc: subprocess.Popen, timeout_s: float = 20.0) -> None: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + if proc.poll() is not None: + raise RuntimeError(f"lightmem_server exited early rc={proc.returncode}") + if redis_ping(host, redis_port) and etcd_health(host, etcd_port): + return + time.sleep(0.1) + raise TimeoutError("timeout waiting for lightmem_server to become ready") + + +@dataclass +class LightMemServerHandle: + proc: subprocess.Popen + kind: str + cmd: list[str] + + def cleanup(self) -> None: + p = self.proc + if p.poll() is not None: + return + + try: + p.send_signal(signal.SIGTERM) + except Exception: + try: + p.terminate() + except Exception: + pass + + try: + p.wait(timeout=8) + except subprocess.TimeoutExpired: + try: + p.kill() + except Exception: + pass + p.wait(timeout=8) + + +def _lightmem_server_command() -> tuple[list[str] | None, str]: + """Return (cmd, kind). Prefer console script, fallback to python -m.""" + exe = which("lightmem_server") + if exe: + return [exe], "lightmem_server" + + # Fallback: run module from repo source tree. + # This still uses the same implementation (server_cli.py) without modifying it. + return [sys.executable, "-m", "light_mem.server_cli"], "python -m light_mem.server_cli" + + +def start_lightmem_server_or_skip(*, host: str, redis_port: int, etcd_port: int, etcd_peer_port: int) -> LightMemServerHandle: + base_cmd, kind = _lightmem_server_command() + if not base_cmd: + print("SKIP: lightmem_server is unavailable.") + raise SystemExit(0) + + cmd = [ + *base_cmd, + "--mode", + "local", + "--index-port", + str(int(redis_port)), + "--coord-port", + str(int(etcd_port)), + "--coord-peer-port", + str(int(etcd_peer_port)), + ] + + env = os.environ.copy() + if kind.startswith("python -m"): + # Ensure repo Python package is importable in the subprocess. + py_src = str(ROOT / "python") + env["PYTHONPATH"] = (py_src + os.pathsep + env.get("PYTHONPATH", "")).rstrip(os.pathsep) + + proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + try: + _wait_lightmem_server_ready(host=host, redis_port=redis_port, etcd_port=etcd_port, proc=proc, timeout_s=25.0) + except Exception: + out = "" + try: + if proc.stdout is not None: + out = proc.stdout.read()[-4000:] + except Exception: + pass + + try: + proc.terminate() + except Exception: + pass + try: + proc.wait(timeout=5) + except Exception: + try: + proc.kill() + except Exception: + pass + + if out: + print("[lightmem_server output tail]\n" + out) + + # Preserve old behavior: skip tests if dependencies are not available. + print("SKIP: failed to start lightmem_server local services (need redis-server and etcd).") + raise SystemExit(0) + + return LightMemServerHandle(proc=proc, kind=kind, cmd=cmd) + + +def redis_resp_command(host: str, port: int, argv: list[str], *, timeout_s: float = 1.0) -> str: + """Send one Redis command over RESP, return first-line reply.""" + payload = "*" + str(len(argv)) + "\r\n" + for a in argv: + b = a.encode("utf-8") + payload += "$" + str(len(b)) + "\r\n" + a + "\r\n" + + with socket.create_connection((host, port), timeout=timeout_s) as s: + s.settimeout(timeout_s) + s.sendall(payload.encode("utf-8")) + + buf = bytearray() + while True: + ch = s.recv(1) + if not ch: + break + buf += ch + if buf.endswith(b"\r\n"): + break + return buf.decode("utf-8", errors="replace").strip() + + +def _redis_read_line(sock: socket.socket) -> bytes: + buf = bytearray() + while True: + ch = sock.recv(1) + if not ch: + raise ConnectionError("redis connection closed") + buf += ch + if buf.endswith(b"\r\n"): + return bytes(buf[:-2]) + + +def _redis_read_exact(sock: socket.socket, n: int) -> bytes: + buf = bytearray() + while len(buf) < n: + chunk = sock.recv(n - len(buf)) + if not chunk: + raise ConnectionError("redis connection closed") + buf += chunk + return bytes(buf) + + +def redis_command(host: str, port: int, argv: list[str], *, timeout_s: float = 1.0) -> str | None: + """Execute a Redis command and return decoded string result. + + Supported replies: + - simple string (+OK) + - error (-ERR ...): raises RuntimeError + - integer (:1) + - bulk string ($len\r\n...): returns str; nil ($-1) -> None + + For multi-node tests we mostly need HGET/HEXISTS/HDEL/DEL. + """ + payload = "*" + str(len(argv)) + "\r\n" + for a in argv: + b = a.encode("utf-8") + payload += "$" + str(len(b)) + "\r\n" + a + "\r\n" + + with socket.create_connection((host, port), timeout=timeout_s) as s: + s.settimeout(timeout_s) + s.sendall(payload.encode("utf-8")) + + first = _redis_read_line(s) + if not first: + return None + + prefix = chr(first[0]) + rest = first[1:] + + if prefix == '+': + return rest.decode("utf-8", errors="replace") + if prefix == '-': + raise RuntimeError(rest.decode("utf-8", errors="replace")) + if prefix == ':': + return rest.decode("utf-8", errors="replace") + if prefix == '$': + try: + ln = int(rest.decode("utf-8", errors="replace")) + except Exception: + return None + if ln < 0: + return None + data = _redis_read_exact(s, ln) + _ = _redis_read_exact(s, 2) # CRLF + return data.decode("utf-8", errors="replace") + + # Not needed for our tests. + raise RuntimeError(f"unsupported RESP reply: {first!r}") + + +def redis_hget_str(host: str, port: int, *, key: str, field: str) -> str | None: + return redis_command(host, port, ["HGET", key, field]) + + +def redis_hexists(host: str, port: int, *, key: str, field: str) -> bool: + v = redis_command(host, port, ["HEXISTS", key, field]) + try: + return int(v or "0") == 1 + except Exception: + return False + + +def redis_hdel(host: str, port: int, *, key: str, field: str) -> None: + _ = redis_command(host, port, ["HDEL", key, field]) + + +def parse_global_mapping(value: str) -> tuple[int, int] | None: + # value like ":" + if not value: + return None + if ":" not in value: + return None + a, b = value.split(":", 1) + try: + return int(a), int(b) + except Exception: + return None + + +@dataclass +class RedisHandle: + host: str + port: int + kind: str + _cleanup_cb: callable | None + + def cleanup(self) -> None: + if self._cleanup_cb is None: + return + cb = self._cleanup_cb + self._cleanup_cb = None + cb() + + +@dataclass +class EtcdHandle: + host: str + port: int + kind: str + _cleanup_cb: callable | None + + def cleanup(self) -> None: + if self._cleanup_cb is None: + return + cb = self._cleanup_cb + self._cleanup_cb = None + cb() + + +def start_redis_local() -> RedisHandle | None: + redis_server = which("redis-server") + if not redis_server: + return None + + port = find_free_port() + tmpdir = Path(tempfile.mkdtemp(prefix="lightmem-test-redis-")) + + cmd = [ + redis_server, + "--bind", + "127.0.0.1", + "--port", + str(port), + "--save", + "", + "--appendonly", + "yes", + "--appendfsync", + "everysec", + "--dir", + str(tmpdir), + ] + + proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + if not wait_port("127.0.0.1", port, timeout_s=10.0): + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + shutil.rmtree(tmpdir, ignore_errors=True) + return None + + def _cleanup() -> None: + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + shutil.rmtree(tmpdir, ignore_errors=True) + + return RedisHandle(host="127.0.0.1", port=port, kind="local", _cleanup_cb=_cleanup) + + +def docker_compose_cmd() -> list[str] | None: + if which("docker"): + try: + run(["docker", "compose", "version"], cwd=ROOT, timeout=10) + return ["docker", "compose"] + except Exception: + pass + if which("docker-compose"): + return ["docker-compose"] + return None + + +def etcd_client(host: str, port: int): + cls = _load_etcd_http_client_cls() + if cls is None: + raise RuntimeError("EtcdV3HttpClient not found in repo sources; cannot create etcd client") + return cls(host=host, port=port) + + +def wait_shard_owners(*, host: str, port: int, prefix: str, num_shards: int, timeout_s: float = 30.0) -> dict[int, str]: + """Wait until every shard has an owner key in etcd.""" + client = etcd_client(host, port) + base = f"{prefix.rstrip('/')}/shards/" + + deadline = time.time() + timeout_s + while time.time() < deadline: + owners: dict[int, str] = {} + # Scan owners. + for value, meta in client.get_prefix(base): + try: + key = meta.key.decode("utf-8") + except Exception: + continue + if not key.endswith("/owner"): + continue + try: + sid_str = key[len(base):].split("/", 1)[0] + sid = int(sid_str) + except Exception: + continue + if 0 <= sid < num_shards and value is not None: + try: + owners[sid] = value.decode("utf-8") + except Exception: + continue + + if len(owners) >= num_shards: + return owners + time.sleep(0.2) + + raise TimeoutError(f"etcd shard owners not ready: have {len(owners)}/{num_shards}") + + +def start_redis_docker_compose_fixed_port() -> RedisHandle | None: + compose = docker_compose_cmd() + if not compose: + return None + + if wait_port("127.0.0.1", 6379, timeout_s=0.2): + return None + + compose_file = ROOT / "docker-compose.redis.yml" + if not compose_file.exists(): + return None + + try: + run([*compose, "-f", str(compose_file), "up", "-d"], cwd=ROOT, timeout=180) + except Exception: + return None + + if not wait_port("127.0.0.1", 6379, timeout_s=20.0): + try: + run([*compose, "-f", str(compose_file), "down", "-v"], cwd=ROOT, timeout=60) + except Exception: + pass + return None + + def _cleanup() -> None: + try: + run([*compose, "-f", str(compose_file), "down", "-v"], cwd=ROOT, timeout=60) + except Exception: + pass + + return RedisHandle(host="127.0.0.1", port=6379, kind="docker-compose", _cleanup_cb=_cleanup) + + +def start_redis_or_skip(*, reuse: bool, host: str, port: int) -> RedisHandle: + if reuse: + return RedisHandle(host=host, port=port, kind="reuse", _cleanup_cb=None) + + h = start_redis_local() + if h: + return h + h = start_redis_docker_compose_fixed_port() + if h: + return h + + print("SKIP: redis-server/docker compose unavailable (or port 6379 busy).") + raise SystemExit(0) + + +def start_etcd_local() -> EtcdHandle | None: + etcd = which("etcd") + if not etcd: + return None + + host = "127.0.0.1" + client_port = find_free_port() + peer_port = find_free_port() + data_dir = Path(tempfile.mkdtemp(prefix="lightmem-test-etcd-")) + + name = "coord" + cmd = [ + etcd, + "--name", + name, + "--data-dir", + str(data_dir), + "--listen-client-urls", + f"http://{host}:{client_port}", + "--advertise-client-urls", + f"http://{host}:{client_port}", + "--listen-peer-urls", + f"http://{host}:{peer_port}", + "--initial-advertise-peer-urls", + f"http://{host}:{peer_port}", + "--initial-cluster", + f"{name}=http://{host}:{peer_port}", + "--initial-cluster-state", + "new", + ] + + proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + if not wait_port(host, client_port, timeout_s=15.0): + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + shutil.rmtree(data_dir, ignore_errors=True) + return None + + def _cleanup() -> None: + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + shutil.rmtree(data_dir, ignore_errors=True) + + return EtcdHandle(host=host, port=client_port, kind="local", _cleanup_cb=_cleanup) + + +def start_etcd_docker() -> EtcdHandle | None: + docker = which("docker") + if not docker: + return None + + host = "127.0.0.1" + client_port = find_free_port() + peer_port = find_free_port() + + # bitnami/etcd image is used elsewhere in repo; keep consistent. + # Use --rm? We want to cleanup deterministically; we'll stop container by name. + name = f"lightmem-test-etcd-{client_port}" + + try: + run( + [ + docker, + "run", + "-d", + "--name", + name, + "-e", + "ALLOW_NONE_AUTHENTICATION=yes", + "-e", + "ETCD_NAME=coord", + "-p", + f"{client_port}:2379", + "-p", + f"{peer_port}:2380", + "bitnami/etcd:3.5", + ], + cwd=ROOT, + timeout=180, + ) + except Exception: + return None + + if not wait_port(host, client_port, timeout_s=20.0): + try: + run([docker, "rm", "-f", name], cwd=ROOT, timeout=30) + except Exception: + pass + return None + + def _cleanup() -> None: + try: + run([docker, "rm", "-f", name], cwd=ROOT, timeout=30) + except Exception: + pass + + return EtcdHandle(host=host, port=client_port, kind="docker", _cleanup_cb=_cleanup) + + +def start_etcd_or_skip(*, reuse: bool, host: str, port: int) -> EtcdHandle: + if reuse: + return EtcdHandle(host=host, port=port, kind="reuse", _cleanup_cb=None) + + h = start_etcd_local() + if h: + return h + h = start_etcd_docker() + if h: + return h + + print("SKIP: etcd/docker unavailable.") + raise SystemExit(0) + + +@dataclass +class ClusterEnv: + redis: RedisHandle + etcd: EtcdHandle + server: LightMemServerHandle | None + storage_dir: Path + cleanup_storage: bool + + def cleanup(self) -> None: + # Stop services first to release file handles. + if self.server is not None: + try: + self.server.cleanup() + finally: + self.server = None + else: + try: + self.etcd.cleanup() + finally: + self.redis.cleanup() + + if self.cleanup_storage: + shutil.rmtree(self.storage_dir, ignore_errors=True) + + +def make_storage_dir(*, base: str | None = None) -> Path: + if base: + p = Path(base) + p.mkdir(parents=True, exist_ok=True) + return p + return Path(tempfile.mkdtemp(prefix="lightmem-multi-node-cache-")) + + +def start_cluster_env(*, reuse_services: bool, redis_host: str, redis_port: int, etcd_host: str, etcd_port: int, + storage_dir: str | None = None, cleanup_storage: bool = True) -> ClusterEnv: + # Multi-node tests now rely on lightmem_server to start/manage both Redis (index) + # and Etcd (coord). Ports are fixed by convention unless explicitly provided. + index_port = int(redis_port) if int(redis_port) > 0 else 6379 + coord_port = int(etcd_port) if int(etcd_port) > 0 else 2379 + + server: LightMemServerHandle | None = None + if not reuse_services: + # lightmem_server --mode local binds to 127.0.0.1. + # Keep behavior explicit to avoid surprising cross-host binds. + if str(redis_host) not in ("127.0.0.1", "localhost") or str(etcd_host) not in ("127.0.0.1", "localhost"): + raise SystemExit("multi_node tests require --redis-host/--etcd-host to be 127.0.0.1 when not using --reuse-services") + server = start_lightmem_server_or_skip(host="127.0.0.1", redis_port=index_port, etcd_port=coord_port, etcd_peer_port=2380) + + r = RedisHandle(host=str(redis_host), port=index_port, kind=("reuse" if reuse_services else "lightmem_server"), _cleanup_cb=None) + e = EtcdHandle(host=str(etcd_host), port=coord_port, kind=("reuse" if reuse_services else "lightmem_server"), _cleanup_cb=None) + sd = make_storage_dir(base=storage_dir) + return ClusterEnv(redis=r, etcd=e, server=server, storage_dir=sd, cleanup_storage=cleanup_storage) + + +def parse_common_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(add_help=False) + p.add_argument("--reuse-services", action="store_true", help="Reuse already running redis/etcd") + p.add_argument("--redis-host", default="127.0.0.1") + p.add_argument("--redis-port", type=int, default=0) + p.add_argument("--etcd-host", default="127.0.0.1") + p.add_argument("--etcd-port", type=int, default=0) + p.add_argument("--storage-dir", default="") + return p.parse_args(argv) + + +def require_ports(ns: argparse.Namespace) -> tuple[str, int, str, int]: + rh = str(ns.redis_host) + rp = int(ns.redis_port) + eh = str(ns.etcd_host) + ep = int(ns.etcd_port) + + if ns.reuse_services: + if rp <= 0 or ep <= 0: + raise SystemExit("--reuse-services requires --redis-port and --etcd-port") + else: + # allow 0 to mean 'auto' for independent mode + if rp < 0 or ep < 0: + raise SystemExit("invalid port") + + return rh, rp, eh, ep + + +def iter_hash_ids(*, count: int, seed: int = 0) -> list[int]: + """Deterministic 128-bit integers for block hashes.""" + out: list[int] = [] + x = (seed & 0xFFFFFFFFFFFFFFFF) | (seed << 64) + for i in range(count): + # A tiny LCG-ish mix into 128-bit space. + x = (x * 6364136223846793005 + 1442695040888963407 + i) & ((1 << 128) - 1) + out.append(x) + return out + + +def build_hash_128s_for_blocks(*, block_hash_ids: Iterable[int], pages_per_block: int) -> list[int]: + """Build a list[int] whose last element of each block is the hash id. + + PyLocalCacheService will take every `pages_per_block`-th element (the last of each block) + as the block hash string. + """ + result: list[int] = [] + dummy = 1 + for hid in block_hash_ids: + for _ in range(pages_per_block - 1): + result.append(dummy) + dummy += 1 + result.append(int(hid)) + return result diff --git a/tests/multi_node/run_all.py b/tests/multi_node/run_all.py new file mode 100755 index 0000000..e877c76 --- /dev/null +++ b/tests/multi_node/run_all.py @@ -0,0 +1,219 @@ +#!/usr/bin/env python3 +"""Run all multi-node tests. + +Behavior: +- Starts Redis + Etcd once +- Runs each test script with --reuse-services and shared storage directory +- Finally stops services + +This mirrors tests/run_all.py but with service reuse. +""" + +from __future__ import annotations + +import argparse +import subprocess +import sys +from pathlib import Path + +from harness import etcd_client, redis_command, require_ports, start_cluster_env + + +def reset_state(*, redis_host: str, redis_port: int, etcd_host: str, etcd_port: int, prefix: str = "lightmem") -> None: + # Redis: these tests treat Redis as an external index; we want a clean DB per test. + try: + redis_command(redis_host, redis_port, ["FLUSHDB"], timeout_s=3.0) + except Exception: + # Best-effort: don't crash run_all on reset failure; the test will surface it. + pass + + # Etcd: delete all keys under prefix. + try: + client = etcd_client(etcd_host, etcd_port) + base = prefix.rstrip("/") + "/" + keys: list[str] = [] + for _value, meta in client.get_prefix(base): + try: + keys.append(meta.key.decode("utf-8")) + except Exception: + continue + for k in keys: + try: + client.delete(k) + except Exception: + pass + except Exception: + pass + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description="Run all LightMem multi_node tests") + p.add_argument("--reuse-services", action="store_true", help="Reuse already running redis/etcd") + p.add_argument("--redis-host", default="127.0.0.1") + p.add_argument("--redis-port", type=int, default=0) + p.add_argument("--etcd-host", default="127.0.0.1") + p.add_argument("--etcd-port", type=int, default=0) + p.add_argument("--storage-dir", default="", help="Optional base storage dir; per-test subdirs will be created") + p.add_argument("--no-reset", action="store_true", help="Do not FLUSHDB / delete etcd prefix between tests") + return p.parse_args(argv) + + +def run_test(test_file: Path, *, redis_host: str, redis_port: int, etcd_host: str, etcd_port: int, storage_dir: str) -> bool: + print(f"\n{'=' * 70}") + print(f"运行 multi_node 测试: {test_file.name}") + print('=' * 70) + + cmd = [ + sys.executable, + str(test_file), + "--reuse-services", + "--redis-host", + redis_host, + "--redis-port", + str(redis_port), + "--etcd-host", + etcd_host, + "--etcd-port", + str(etcd_port), + "--storage-dir", + storage_dir, + ] + + try: + result = subprocess.run( + cmd, + cwd=str(test_file.parent), + timeout=600, + capture_output=True, + text=True, + ) + + out = (result.stdout or "") + (result.stderr or "") + if out.strip(): + print(out.rstrip()) + + if result.returncode == 0: + return True + + if result.returncode < 0: + print(f"✗ 测试被信号终止: {test_file.name} (signal={-result.returncode})") + else: + print(f"✗ 测试失败: {test_file.name} (code={result.returncode})") + + # Print a small tail to make failures actionable. + tail = out[-4000:] if out else "" + if tail and tail != out: + print("[output tail]\n" + tail) + return False + except subprocess.TimeoutExpired: + print(f"✗ 测试超时: {test_file.name}") + return False + + +def main() -> int: + tests_dir = Path(__file__).parent + + ns = parse_args(sys.argv[1:]) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + test_files = [ + tests_dir / "test_00_bootstrap_empty_dir.py", + tests_dir / "test_01_restart_recover_existing_dir.py", + tests_dir / "test_02_assignment_balance_and_uniqueness.py", + tests_dir / "test_03_write_goes_to_owned_shards.py", + tests_dir / "test_04_cross_node_dedupe_same_hash.py", + tests_dir / "test_05_join_rebalance_minimal_movement.py", + tests_dir / "test_06_leave_lease_expire_reassign.py", + tests_dir / "test_07_redis_flush_and_recover_to_redis.py", + tests_dir / "test_08_cross_node_read_lru_window.py", + tests_dir / "test_09_crash_recovery_subprocess.py", + tests_dir / "test_10_watch_prefix_callback.py", + tests_dir / "test_11_crc_validation.py", + tests_dir / "test_12_redis_recovery.py", + tests_dir / "test_13_concurrent_init_shared_dir.py", + ] + + print("LightMem multi_node 测试套件") + print(f"测试目录: {tests_dir}") + print(f"总测试数: {len(test_files)}") + + # Use port=0 (auto) for local binaries; for reuse, harness ignores. + storage_dir = str(getattr(ns, "storage_dir", "") or "").strip() or None + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port if bool(ns.reuse_services) else 6379, + etcd_host=etcd_host, + etcd_port=etcd_port if bool(ns.reuse_services) else 2379, + storage_dir=storage_dir, + cleanup_storage=storage_dir is None, + ) + + mode = "REUSE" if bool(ns.reuse_services) else "START" + reset = "OFF" if bool(ns.no_reset) else "ON" + print(f"\n[run_all] services mode={mode} reset_between_tests={reset}") + print(f"[run_all] redis={env.redis.host}:{env.redis.port} kind={env.redis.kind}") + print(f"[run_all] etcd={env.etcd.host}:{env.etcd.port} kind={env.etcd.kind}") + print(f"[run_all] storage_root={env.storage_dir}") + + try: + results: dict[str, bool] = {} + passed = 0 + failed = 0 + + for tf in test_files: + if not tf.exists(): + print(f"⚠ 跳过不存在的测试: {tf.name}") + continue + + if not bool(ns.no_reset): + print(f"\n[run_all] reset redis/etcd state for {tf.name}") + reset_state( + redis_host=env.redis.host, + redis_port=env.redis.port, + etcd_host=env.etcd.host, + etcd_port=env.etcd.port, + prefix="lightmem", + ) + + # Isolate disk state per test to avoid cross-test interference while + # still keeping the "shared directory across nodes" property within a test. + test_storage_dir = env.storage_dir / tf.stem + if test_storage_dir.exists(): + # Clean up stale files from a previous run. + import shutil + + shutil.rmtree(test_storage_dir, ignore_errors=True) + test_storage_dir.mkdir(parents=True, exist_ok=True) + + ok = run_test( + tf, + redis_host=env.redis.host, + redis_port=env.redis.port, + etcd_host=env.etcd.host, + etcd_port=env.etcd.port, + storage_dir=str(test_storage_dir), + ) + results[tf.name] = ok + if ok: + passed += 1 + else: + failed += 1 + + print("\n" + "=" * 70) + print("multi_node 测试总结") + print("=" * 70) + for name, ok in results.items(): + status = "✓ 通过" if ok else "✗ 失败" + print(f"{status:8} {name}") + print("=" * 70) + print(f"总计: {passed} 通过, {failed} 失败, 共 {passed + failed} 个测试") + print("=" * 70) + + return 0 if failed == 0 else 1 + finally: + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/stability_dynamic_nodes.py b/tests/multi_node/stability_dynamic_nodes.py new file mode 100755 index 0000000..9abb26a --- /dev/null +++ b/tests/multi_node/stability_dynamic_nodes.py @@ -0,0 +1,977 @@ +#!/usr/bin/env python3 +"""共享存储多节点稳定性测试（支持单节点） + +目标（符合 multi_node 的真实语义）： +- 启动 N 个节点（N 可配置，默认 1），指向同一个 storage 前缀（同一份缓存文件/目录，不是每节点一份）。 +- 通过 Etcd 动态分配 shard owner（写入权限），通过 Redis 维护索引/映射。 +- 持续 RW 写满触发 LRU 淘汰。 +- 当 N>1 时，跨节点读验证：节点 A 写入的 key，节点 B 必须能读到且校验字节一致（反之亦然）。 + +实现方式： +- 本脚本仅做“编排器”，通过 subprocess 启动两个 tests/multi_node/worker_node_ops.py 进程。 +- worker 使用 "loop_rw_verify"：持续写入唯一 key + 刷新 hot set + probe-file 交叉验证。 + +注意： +- 不在同一进程内创建多个 PyLocalCacheService（避免潜在崩溃/线程问题）。 +""" + +from __future__ import annotations + +import argparse +import json +import os +import random +import signal +import subprocess +import sys +import tempfile +import time +import uuid +from pathlib import Path + + +HERE = Path(__file__).resolve() +REPO_ROOT = HERE.parents[2] +WORKER = REPO_ROOT / "tests" / "multi_node" / "worker_node_ops.py" + + +def _gb_to_bytes(gb: float) -> int: + return int(float(gb) * 1024 * 1024 * 1024) + + +def _wait_file(path: Path, timeout_s: float) -> None: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + if path.exists(): + return + time.sleep(0.02) + raise TimeoutError(f"timeout waiting for file: {path}") + + +def _read_json(path: Path) -> dict: + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return {} + + +def _tail_text(path: Path, max_bytes: int = 24 * 1024) -> str: + try: + data = path.read_bytes() + if len(data) > max_bytes: + data = data[-max_bytes:] + return data.decode("utf-8", errors="replace") + except Exception: + return "" + + +def _split_pipeline(total_sec: float, write_sec: float, read_sec: float, verify_sec: float) -> tuple[float, float, float]: + total = max(1.0, float(total_sec)) + if write_sec > 0 or read_sec > 0 or verify_sec > 0: + w = max(0.0, float(write_sec)) + r = max(0.0, float(read_sec)) + v = max(0.0, float(verify_sec)) + if w + r + v <= 0: + return total, 0.0, 0.0 + # If user-specified sum differs, scale to fit total. + scale = total / (w + r + v) + return w * scale, r * scale, v * scale + + # Auto split: bias toward write+read bandwidth, keep small verify. + v = max(5.0, total * 0.05) + r = max(10.0, total * 0.25) + w = max(10.0, total - r - v) + # If total is tiny, clamp. + if w + r + v > total: + # shrink in order r then v. + overflow = (w + r + v) - total + r = max(0.0, r - overflow) + overflow2 = (w + r + v) - total + if overflow2 > 0: + v = max(0.0, v - overflow2) + return float(w), float(r), float(v) + + +def _spawn_worker_with_overrides( + *, + node_id: str, + storage_prefix: Path, + run_dir: Path, + args: argparse.Namespace, + op: str, + duration_sec: float, + probe_file: str, + hash_id_hex: str, +) -> subprocess.Popen: + # Clone args shallowly by creating a tiny shim namespace. + tmp = argparse.Namespace(**vars(args)) + tmp.mode = "verify" if op == "loop_rw_verify" else "max_bw" + tmp.duration_sec = float(duration_sec) + tmp.probe_file = str(probe_file) + + # Build command similarly to _spawn_worker but with explicit op. + started_file = run_dir / f"{node_id}.started" + ready_file = run_dir / f"{node_id}.ready" + done_file = run_dir / f"{node_id}.done" + stats_file = run_dir / f"{node_id}.stats.json" + log_file = run_dir / f"{node_id}.log" + + base_hex = str(hash_id_hex) + + cmd: list[str] = [ + sys.executable, + "-u", + str(WORKER), + "--storage-dir", + str(storage_prefix), + "--storage-size", + str(int(tmp.storage_size_bytes)), + "--num-shard", + str(int(tmp.num_shard)), + "--num-worker", + str(int(tmp.num_worker)), + # Coordination endpoints are added below (unless disabled). + "--node-id", + str(node_id), + "--ttl", + str(int(tmp.ttl)), + "--num-pages", + str(int(tmp.num_pages)), + "--page-bytes", + str(int(tmp.page_bytes)), + "--op", + str(op), + "--hash-id", + str(base_hex), + "--duration-sec", + str(float(tmp.duration_sec)), + "--report-interval-sec", + str(float(tmp.report_interval_sec)), + "--stats-file", + str(stats_file), + "--started-file", + str(started_file), + "--ready-file", + str(ready_file), + "--done-file", + str(done_file), + "--resident-window", + str(int(tmp.resident_window)), + "--evict-probe-gap", + str(int(tmp.evict_probe_gap)), + "--debug-dump-file", + str(run_dir / f"{node_id}.debug.json"), + ] + + if bool(getattr(tmp, "disable_coord", False)): + cmd.append("--disable-coord") + else: + cmd.extend(["--redis", str(tmp.redis), "--etcd", str(tmp.etcd)]) + + if op == "bench_write": + cmd.extend(["--batch-blocks", str(int(tmp.batch_blocks))]) + if op == "bench_read": + cmd.extend( + [ + "--batch-blocks", + str(int(tmp.batch_blocks)), + "--read-window-blocks", + str(int(tmp.read_window_blocks)), + ] + ) + + if probe_file: + cmd.extend( + [ + "--probe-file", + str(probe_file), + "--hot-set-size", + str(int(tmp.hot_set_size)), + "--probe-publish-interval-sec", + str(float(tmp.probe_publish_interval_sec)), + "--probe-read-interval-sec", + str(float(tmp.probe_read_interval_sec)), + "--probe-ttl-sec", + str(float(tmp.probe_ttl_sec)), + "--probe-max-read", + str(int(tmp.probe_max_read)), + ] + ) + + env = os.environ.copy() + # etcd-mode is inferred from PyLocalCacheService(coord_endpoints=...). + + # Optional: isolate etcd/redis namespaces to avoid interference with other live nodes/runs. + if not bool(getattr(tmp, "disable_coord", False)): + index_prefix = str(getattr(args, "index_prefix", "") or "").strip() + if index_prefix: + cmd.extend(["--index-prefix", index_prefix]) + + if bool(getattr(args, "expect_own_all_shards", False)): + env.setdefault("LIGHTMEM_EXPECT_OWN_ALL_SHARDS", "1") + + mode = str(getattr(args, "worker_stdout", "file")) + if mode == "inherit": + proc = subprocess.Popen(cmd, env=env) + proc._lightmem_log_fh = None # type: ignore[attr-defined] + else: + log_fh = log_file.open("w", encoding="utf-8") + proc = subprocess.Popen(cmd, stdout=log_fh, stderr=subprocess.STDOUT, env=env) + proc._lightmem_log_fh = log_fh # type: ignore[attr-defined] + return proc + + +def _spawn_worker(*, node_id: str, storage_prefix: Path, run_dir: Path, args: argparse.Namespace) -> subprocess.Popen: + started_file = run_dir / f"{node_id}.started" + ready_file = run_dir / f"{node_id}.ready" + done_file = run_dir / f"{node_id}.done" + stats_file = run_dir / f"{node_id}.stats.json" + log_file = run_dir / f"{node_id}.log" + + # Use a safe int64 decimal base id (some backends assume stoll-able values). + # Keep it in signed int64 range. + base = random.getrandbits(63) + base_hex = str(base) + + op = "loop_rw_verify" if str(args.mode) == "verify" else "bench_write" + + cmd: list[str] = [ + sys.executable, + "-u", + str(WORKER), + "--storage-dir", + str(storage_prefix), + "--storage-size", + str(int(args.storage_size_bytes)), + "--num-shard", + str(int(args.num_shard)), + "--num-worker", + str(int(args.num_worker)), + # Coordination endpoints are added below (unless disabled). + "--node-id", + str(node_id), + "--ttl", + str(int(args.ttl)), + "--num-pages", + str(int(args.num_pages)), + "--page-bytes", + str(int(args.page_bytes)), + "--op", + str(op), + "--hash-id", + str(base_hex), + "--duration-sec", + str(float(args.duration_sec)), + "--report-interval-sec", + str(float(args.report_interval_sec)), + "--stats-file", + str(stats_file), + "--started-file", + str(started_file), + "--ready-file", + str(ready_file), + "--done-file", + str(done_file), + "--resident-window", + str(int(args.resident_window)), + "--evict-probe-gap", + str(int(args.evict_probe_gap)), + "--debug-dump-file", + str(run_dir / f"{node_id}.debug.json"), + ] + + if bool(getattr(args, "disable_coord", False)): + cmd.append("--disable-coord") + else: + cmd.extend(["--redis", str(args.redis), "--etcd", str(args.etcd)]) + + # Cross-node verification is only enabled when multiple nodes run. + if getattr(args, "probe_file", ""): + cmd.extend( + [ + "--probe-file", + str(args.probe_file), + "--hot-set-size", + str(int(args.hot_set_size)), + "--probe-publish-interval-sec", + str(float(args.probe_publish_interval_sec)), + "--probe-read-interval-sec", + str(float(args.probe_read_interval_sec)), + "--probe-ttl-sec", + str(float(args.probe_ttl_sec)), + "--probe-max-read", + str(int(args.probe_max_read)), + ] + ) + + # High-throughput bench params. + if str(args.mode) == "max_bw": + cmd.extend(["--batch-blocks", str(int(args.batch_blocks))]) + + env = os.environ.copy() + # etcd-mode is inferred from PyLocalCacheService(coord_endpoints=...). + + # Optional: isolate etcd/redis namespaces to avoid interference with other live nodes/runs. + if not bool(getattr(args, "disable_coord", False)): + index_prefix = str(getattr(args, "index_prefix", "") or "").strip() + if index_prefix: + cmd.extend(["--index-prefix", index_prefix]) + + if bool(getattr(args, "expect_own_all_shards", False)): + env.setdefault("LIGHTMEM_EXPECT_OWN_ALL_SHARDS", "1") + + mode = str(getattr(args, "worker_stdout", "file")) + if mode == "inherit": + proc = subprocess.Popen(cmd, env=env) + proc._lightmem_log_fh = None # type: ignore[attr-defined] + else: + log_fh = log_file.open("w", encoding="utf-8") + proc = subprocess.Popen(cmd, stdout=log_fh, stderr=subprocess.STDOUT, env=env) + proc._lightmem_log_fh = log_fh # type: ignore[attr-defined] + return proc + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description="共享存储（Redis+Etcd）RW + LRU（可选跨节点读验证）") + p.add_argument("--num-nodes", type=int, default=1, help="节点数量（默认 1；>1 时启用跨节点读验证）") + p.add_argument( + "--node-id", + default="", + help=( + "可选：固定 node_id（用于多机各跑一份脚本时保持稳定身份）。" + "留空则每次运行自动生成，避免两台机器 node_id 冲突。" + ), + ) + p.add_argument( + "--mode", + choices=["verify", "max_bw", "pipeline"], + default="verify", + help="verify: 正确性/跨节点验证(吞吐较低)；max_bw: 高吞吐批量写入；pipeline: 写带宽→读带宽→短验证", + ) + p.add_argument("--storage-dir", required=True, help="共享缓存目录（两个节点将共享同一份 shard 文件前缀）") + p.add_argument("--storage-size-gb", type=float, default=120.0, help="共享存储容量（GB），默认 40GB") + p.add_argument("--num-shard", type=int, default=192) + p.add_argument("--num-worker", type=int, default=32) + + p.add_argument( + "--disable-coord", + action="store_true", + help="禁用 Redis/Etcd 通信，仅测试本地存储读写带宽（安全起见仅支持 --num-nodes 1）", + ) + p.add_argument("--redis", default="", help="host:port") + p.add_argument("--etcd", default="", help="host:port") + p.add_argument( + "--index-prefix", + default="", + help="索引/协调前缀（同时作用于 etcd key 与 Redis key；为空则每次运行自动生成以隔离测试）", + ) + p.add_argument( + "--expect-own-all-shards", + action="store_true", + help=( + "严格单节点模式：要求本节点必须拥有全部 shard 才开始读写。" + "当你要在不同服务器同时跑同一测试（共享同一 --index-prefix）时不要开启。" + ), + ) + p.add_argument("--ttl", type=int, default=6) + p.add_argument("--reconcile-sec", type=float, default=1.0) + + p.add_argument("--duration-sec", type=float, default=120.0) + p.add_argument("--report-interval-sec", type=float, default=2.0) + + # pipeline phase durations; 0 means auto-split from duration-sec. + p.add_argument("--write-sec", type=float, default=0.0, help="pipeline: 写带宽阶段时长(秒)，0=自动") + p.add_argument("--read-sec", type=float, default=0.0, help="pipeline: 读带宽阶段时长(秒)，0=自动") + p.add_argument("--verify-sec", type=float, default=0.0, help="pipeline: 验证阶段时长(秒)，0=自动") + + p.add_argument("--page-bytes", type=int, default=1024 * 1024, help="每页字节数，默认 1024KB") + p.add_argument("--num-pages", type=int, default=256, help="kvcache 页数；越大越吃内存（num_pages*page_bytes）") + + p.add_argument("--resident-window", type=int, default=64) + p.add_argument("--hot-set-size", type=int, default=8) + p.add_argument("--evict-probe-gap", type=int, default=1024) + p.add_argument("--batch-blocks", type=int, default=64, help="max_bw 模式：每次 create 写入的 blocks 数") + p.add_argument("--read-window-blocks", type=int, default=2048, help="pipeline/max_bw: 读带宽阶段读取的热窗口 blocks 数") + + p.add_argument("--probe-publish-interval-sec", type=float, default=1.0) + p.add_argument("--probe-read-interval-sec", type=float, default=1.0) + p.add_argument("--probe-ttl-sec", type=float, default=30.0) + p.add_argument("--probe-max-read", type=int, default=8) + + p.add_argument("--run-dir", default="", help="可选：保存日志/统计的目录；默认使用 /tmp 下临时目录") + p.add_argument( + "--worker-stdout", + choices=["file", "inherit"], + default="", + help="worker 输出方式：file=重定向到 run_dir/*.log；inherit=直接输出到终端（单节点建议）", + ) + args = p.parse_args(argv) + + num_nodes = max(1, int(args.num_nodes)) + + if bool(getattr(args, "expect_own_all_shards", False)) and num_nodes != 1: + raise ValueError("--expect-own-all-shards 仅适用于 --num-nodes 1") + + if bool(getattr(args, "disable_coord", False)): + if num_nodes != 1: + raise ValueError("--disable-coord 仅支持 --num-nodes 1（多进程共享目录会导致数据/索引不一致）") + else: + if not str(args.redis).strip() or not str(args.etcd).strip(): + raise ValueError("--redis/--etcd 不能为空（或使用 --disable-coord 测纯带宽）") + + # Default worker stdout behavior. + if not str(getattr(args, "worker_stdout", "")): + # Default to file so only this orchestrator prints to the console. + # Use --worker-stdout inherit when debugging a single worker. + args.worker_stdout = "file" + + storage_dir = Path(str(args.storage_dir)).resolve() + storage_dir.mkdir(parents=True, exist_ok=True) + + args.storage_size_bytes = _gb_to_bytes(float(args.storage_size_gb)) # type: ignore[attr-defined] + + if args.run_dir: + run_dir = Path(args.run_dir).resolve() + run_dir.mkdir(parents=True, exist_ok=True) + else: + run_dir = Path(tempfile.mkdtemp(prefix="lightmem_stability_dynamic_nodes_")) + + # Only enable cross-node probes when multiple nodes are present. + if num_nodes > 1: + args.probe_file = str(run_dir / "probe.jsonl") # type: ignore[attr-defined] + else: + args.probe_file = "" # type: ignore[attr-defined] + + shard_prefix = storage_dir / "shard" + + print("=" * 80) + print("共享存储稳定性测试（Redis+Etcd）") + print("=" * 80) + print("关键信息:") + print(f" num_nodes: {num_nodes}") + print(f" storage_dir: {storage_dir}") + print(f" shard_prefix: {shard_prefix}") + print(f" storage_size: {float(args.storage_size_gb):.2f} GB") + print(f" num_shard: {int(args.num_shard)}") + if bool(getattr(args, "disable_coord", False)): + print(" coord: DISABLED (bandwidth-only)") + else: + print(f" redis: {args.redis}") + print(f" etcd: {args.etcd}") + print(f" duration: {float(args.duration_sec):.1f} s") + print(f" page_bytes: {int(args.page_bytes)}") + print(f" num_pages: {int(args.num_pages)}") + print(f" run_dir(日志/统计): {run_dir}") + if bool(getattr(args, "expect_own_all_shards", False)): + print(" ownership_gate: expect_own_all_shards=1") + print() + + if not WORKER.exists(): + raise FileNotFoundError(f"worker not found: {WORKER}") + + run_id = uuid.uuid4().hex[:8] + fixed = str(getattr(args, "node_id", "") or "").strip() + if fixed and num_nodes == 1: + node_ids = [fixed] + elif fixed and num_nodes > 1: + node_ids = [f"{fixed}-{i:02d}" for i in range(num_nodes)] + else: + # Default: unique per-run IDs to avoid collisions across machines. + node_ids = [f"stability-{i:02d}-{run_id}" for i in range(num_nodes)] + # NOTE: Some Redis/CPP paths parse ids via stoll (signed 64-bit). + # Keep ids within int64 range to avoid invalid_argument/overflow. + node_base_hex = {nid: str(random.randint(1, (1 << 62) - 1)) for nid in node_ids} + + if not bool(getattr(args, "disable_coord", False)): + # If user didn't specify prefixes, default to an isolated namespace per run. + # Otherwise, a single-node test can end up sharing shard ownership with other live nodes + # registered under the default "lightmem" prefix, causing early per-shard eviction. + index_prefix_arg = (str(getattr(args, "index_prefix", "")) or "").strip() + + index_prefix = index_prefix_arg or f"lightmem_test_{run_id}" + args.index_prefix = index_prefix # type: ignore[attr-defined] + + print(f" index_prefix: {index_prefix}") + + procs: list[subprocess.Popen] = [] + + def _terminate_all(): + for pr in procs: + try: + pr.terminate() + except Exception: + pass + + def _kill_all(): + for pr in procs: + try: + pr.kill() + except Exception: + pass + + def _sigint(_signum, _frame): + _terminate_all() + raise SystemExit(130) + + signal.signal(signal.SIGINT, _sigint) + signal.signal(signal.SIGTERM, _sigint) + + def _run_phase(*, label: str, op: str, duration_sec: float, enable_probe: bool, per_node_hash_hex: dict[str, str]) -> dict: + nonlocal procs + procs = [] + # Reset marker files from previous phase. + for nid in node_ids: + for suf in ("started", "ready", "done"): + pth = run_dir / f"{nid}.{suf}" + try: + if pth.exists(): + pth.unlink() + except Exception: + pass + + # Also reset per-node stats/debug logs so progress output never uses stale data. + for pth in ( + run_dir / f"{nid}.stats.json", + run_dir / f"{nid}.debug.json", + run_dir / f"{nid}.log", + ): + try: + if pth.exists(): + pth.unlink() + except Exception: + pass + + probe_file = str(run_dir / "probe.jsonl") if (enable_probe and num_nodes > 1) else "" + if probe_file: + try: + Path(probe_file).unlink(missing_ok=True) + except Exception: + pass + + print("=" * 80) + print(f"Phase: {label} (op={op}, duration={duration_sec:.1f}s)") + print("=" * 80) + + for nid in node_ids: + procs.append( + _spawn_worker_with_overrides( + node_id=nid, + storage_prefix=shard_prefix, + run_dir=run_dir, + args=args, + op=op, + duration_sec=float(duration_sec), + probe_file=probe_file, + hash_id_hex=str(per_node_hash_hex[nid]), + ) + ) + + for nid in node_ids: + _wait_file(run_dir / f"{nid}.started", timeout_s=60.0) + + ready_timeout = 120.0 if num_nodes > 1 else 90.0 + for nid in node_ids: + _wait_file(run_dir / f"{nid}.ready", timeout_s=ready_timeout) + + print(f"{num_nodes} 个节点已就绪，开始运行 {label}…") + + t0 = time.time() + next_report = time.time() + float(args.report_interval_sec) + + last_reads_ok: dict[str, int] = {nid: 0 for nid in node_ids} + per_node_delta_reads: dict[str, int] = {nid: 0 for nid in node_ids} + + while True: + now = time.time() + rcs = [pr.poll() for pr in procs] + + if now >= next_report: + total_thr = 0.0 + total_writes = 0 + total_writes_hot = 0 + total_reads_ok = 0 + total_reads_miss = 0 + total_probe_ok = 0 + total_probe_fail = 0 + eviction_any_live = False + eviction_cnt_total = 0 + + per_node: dict[str, dict] = {} + for nid in node_ids: + s = _read_json(run_dir / f"{nid}.stats.json") + per_node[nid] = s + total_thr += float(s.get("inst_throughput_mb_s", s.get("throughput_mb_s", 0.0)) or 0.0) + total_writes += int(s.get("writes", 0) or 0) + total_writes_hot += int(s.get("writes_hot", 0) or 0) + total_reads_ok += int(s.get("reads_ok", 0) or 0) + total_reads_miss += int(s.get("reads_miss", 0) or 0) + total_probe_ok += int(s.get("probe_reads_ok", 0) or 0) + total_probe_fail += int(s.get("probe_reads_fail", 0) or 0) + eviction_any_live = eviction_any_live or bool(s.get("eviction_observed", False)) + eviction_cnt_total += int(s.get("eviction_count", 0) or 0) + + elapsed = max(0.0, now - t0) + msg = f"[{elapsed:6.1f}s] 合计吞吐={total_thr:8.1f} MB/s" + if op in ("bench_write", "loop_rw_verify"): + msg += f" 写入={total_writes:7d}" + # Cumulative written capacity (GB). + try: + block_size = int(_read_json(run_dir / f"{node_ids[0]}.stats.json").get("block_size", 0) or 0) + except Exception: + block_size = 0 + if block_size > 0: + total_written_bytes = int(total_writes + total_writes_hot) * int(block_size) + total_written_gb = float(total_written_bytes) / (1024.0**3) + msg += f" 写入≈{total_written_gb:6.2f}GB" + + # Effective capacity visibility (bench_write only). + if op == "bench_write": + eff_bytes = 0 + owned_list: list[int] = [] + for nid in node_ids: + s = per_node.get(nid, {}) + cap_eff = int(s.get("capacity_blocks_effective", 0) or 0) + bs = int(s.get("block_size", 0) or 0) + if cap_eff > 0 and bs > 0: + eff_bytes += int(cap_eff) * int(bs) + owned_list.append(int(s.get("owned_shards", 0) or 0)) + if eff_bytes > 0: + msg += f" cap_eff≈{(float(eff_bytes) / (1024.0**3)):6.2f}GB" + if owned_list: + msg += f" owned_shards=[{min(owned_list)}-{max(owned_list)}]" + if op in ("bench_read", "loop_rw_verify"): + msg += f" 读OK={total_reads_ok:7d}" + # Per-interval read capacity (GB), not cumulative. + try: + block_size = int(_read_json(run_dir / f"{node_ids[0]}.stats.json").get("block_size", 0) or 0) + except Exception: + block_size = 0 + if block_size > 0: + delta_reads = 0 + for nid in node_ids: + s = _read_json(run_dir / f"{nid}.stats.json") + cur = int(s.get("reads_ok", 0) or 0) + prev = int(last_reads_ok.get(nid, 0)) + if cur >= prev: + d = (cur - prev) + per_node_delta_reads[nid] = d + delta_reads += d + else: + per_node_delta_reads[nid] = 0 + last_reads_ok[nid] = cur + delta_gb = float(int(delta_reads) * int(block_size)) / (1024.0**3) + msg += f" 读≈{delta_gb:6.2f}GB" + if enable_probe and num_nodes > 1: + msg += f" probe_ok={total_probe_ok:5d} probe_fail={total_probe_fail:5d}" + msg += f" evict={int(eviction_any_live)} evict_cnt={int(eviction_cnt_total)}" + print(msg) + + # Per-node breakdown for multi-node visibility. + if num_nodes > 1: + for nid in node_ids: + s = per_node.get(nid, {}) + ev = int(bool(s.get("eviction_observed", False))) + evc = int(s.get("eviction_count", 0) or 0) + block_size = int(s.get("block_size", 0) or 0) + + if op == "bench_read": + thr = float(s.get("inst_throughput_mb_s", s.get("throughput_mb_s", 0.0)) or 0.0) + rok = int(s.get("reads_ok", 0) or 0) + d = int(per_node_delta_reads.get(nid, 0) or 0) + line = f" - {nid}: thr={thr:8.1f}MB/s 读OK={rok:7d}" + if block_size > 0: + node_read_gb = float(int(d) * int(block_size)) / (1024.0**3) + line += f" 读≈{node_read_gb:6.2f}GB" + line += f" evict={ev}" + line += f" evict_cnt={evc}" + else: + w = int(s.get("writes", 0) or 0) + wh = int(s.get("writes_hot", 0) or 0) + line = f" - {nid}: 写入={w:7d}" + if block_size > 0: + node_written_bytes = int(w + wh) * int(block_size) + node_written_gb = float(node_written_bytes) / (1024.0**3) + line += f" 写入≈{node_written_gb:6.2f}GB" + line += f" evict={ev}" + line += f" evict_cnt={evc}" + print(line) + next_report = now + float(args.report_interval_sec) + + if all(rc is not None for rc in rcs): + break + + if now - t0 > float(duration_sec) + 120.0: + raise TimeoutError(f"phase timeout: {label}") + + time.sleep(0.1) + + # Check exit codes. + bad: list[tuple[str, int | None]] = [] + for nid, pr in zip(node_ids, procs, strict=True): + if pr.returncode != 0: + bad.append((nid, pr.returncode)) + if bad: + print("\n节点进程异常退出：") + for nid, rc in bad: + print(f" {nid} returncode={rc} log={run_dir / f'{nid}.log'}") + for nid, _ in bad[:2]: + print(f"\n--- {nid} log tail ---") + print(_tail_text(run_dir / f"{nid}.log")) + raise RuntimeError(f"phase failed: {label}") + + return {nid: _read_json(run_dir / f"{nid}.stats.json") for nid in node_ids} + + try: + if str(args.mode) == "pipeline": + wsec, rsec, vsec = _split_pipeline(float(args.duration_sec), float(args.write_sec), float(args.read_sec), float(args.verify_sec)) + + # Phase 1: write bandwidth. + write_stats = _run_phase(label="write_bw", op="bench_write", duration_sec=wsec, enable_probe=False, per_node_hash_hex=node_base_hex) + eviction_any_write = any(bool(s.get("eviction_observed", False)) for s in write_stats.values()) + + # Read from the most recent window to avoid evicted early keys. + window = max(1, int(args.read_window_blocks)) + per_node_read_hex: dict[str, str] = {} + for nid in node_ids: + s = write_stats.get(nid, {}) + base_int = int(node_base_hex[nid]) + try: + last_int = int(s.get("last_written_id", base_int)) + except Exception: + last_int = base_int + read_base = max(0, last_int - window + 1) + per_node_read_hex[nid] = str(int(read_base)) + + # Phase 2: read bandwidth (best-effort over a recent hot window). + _ = _run_phase(label="read_bw", op="bench_read", duration_sec=rsec, enable_probe=False, per_node_hash_hex=per_node_read_hex) + + # Phase 3: short correctness verification. + _ = _run_phase(label="verify", op="loop_rw_verify", duration_sec=vsec, enable_probe=True, per_node_hash_hex=node_base_hex) + + # Final stats from last phase. + final_stats = {nid: _read_json(run_dir / f"{nid}.stats.json") for nid in node_ids} + else: + # Single-phase modes (backward-compatible): verify or max_bw. + for nid in node_ids: + procs.append(_spawn_worker(node_id=nid, storage_prefix=shard_prefix, run_dir=run_dir, args=args)) + + for nid in node_ids: + _wait_file(run_dir / f"{nid}.started", timeout_s=60.0) + + ready_timeout = 120.0 if num_nodes > 1 else 90.0 + for nid in node_ids: + _wait_file(run_dir / f"{nid}.ready", timeout_s=ready_timeout) + + if num_nodes > 1: + print(f"{num_nodes} 个节点已就绪，开始持续 RW + LRU + 跨节点读验证…") + else: + if str(args.mode) == "max_bw": + print("单节点已就绪，开始高吞吐批量写入 + LRU…") + else: + print("单节点已就绪，开始持续 RW + LRU…") + + t0 = time.time() + next_report = time.time() + float(args.report_interval_sec) + + # For per-interval read volume (non-cumulative). + last_reads_ok: dict[str, int] = {nid: 0 for nid in node_ids} + per_node_delta_reads: dict[str, int] = {nid: 0 for nid in node_ids} + + while True: + now = time.time() + + rcs = [pr.poll() for pr in procs] + + if now >= next_report: + total_thr = 0.0 + total_writes = 0 + total_writes_hot = 0 + total_probe_ok = 0 + total_probe_fail = 0 + eviction_any_live = False + eviction_cnt_total = 0 + + per_node: dict[str, dict] = {} + + for nid in node_ids: + s = _read_json(run_dir / f"{nid}.stats.json") + per_node[nid] = s + total_thr += float(s.get("inst_throughput_mb_s", s.get("throughput_mb_s", 0.0)) or 0.0) + total_writes += int(s.get("writes", 0) or 0) + total_writes_hot += int(s.get("writes_hot", 0) or 0) + total_probe_ok += int(s.get("probe_reads_ok", 0) or 0) + total_probe_fail += int(s.get("probe_reads_fail", 0) or 0) + eviction_any_live = eviction_any_live or bool(s.get("eviction_observed", False)) + eviction_cnt_total += int(s.get("eviction_count", 0) or 0) + elapsed = max(0.0, now - t0) + + # Capacity reporting uses block_size from any available stats. + try: + sample = _read_json(run_dir / f"{node_ids[0]}.stats.json") + block_size = int(sample.get("block_size", 0) or 0) + op_name = str(sample.get("op", "") or "") + except Exception: + block_size = 0 + op_name = "" + write_gb = None + if block_size > 0 and op_name in ("bench_write", "loop_rw_verify"): + write_gb = float(int(total_writes + total_writes_hot) * int(block_size)) / (1024.0**3) + + # Read GB per interval (best-effort). + read_gb = None + if block_size > 0 and op_name in ("bench_read", "loop_rw_verify"): + delta_reads = 0 + for nid in node_ids: + s = _read_json(run_dir / f"{nid}.stats.json") + cur = int(s.get("reads_ok", 0) or 0) + prev = int(last_reads_ok.get(nid, 0)) + if cur >= prev: + d = (cur - prev) + per_node_delta_reads[nid] = d + delta_reads += d + else: + per_node_delta_reads[nid] = 0 + last_reads_ok[nid] = cur + read_gb = float(int(delta_reads) * int(block_size)) / (1024.0**3) + + if num_nodes > 1: + msg = ( + f"[{elapsed:6.1f}s] " + f"节点数={num_nodes} 合计吞吐={total_thr:8.1f} MB/s 写入={total_writes:7d}" + ) + if write_gb is not None: + msg += f" 写入≈{write_gb:6.2f}GB" + if read_gb is not None: + msg += f" 读≈{read_gb:6.2f}GB" + msg += f" probe_ok={total_probe_ok:5d} probe_fail={total_probe_fail:5d} evict={int(eviction_any_live)}" + msg += f" evict_cnt={int(eviction_cnt_total)}" + print(msg) + + for nid in node_ids: + s = per_node.get(nid, {}) + blk = int(s.get("block_size", 0) or 0) + ev = int(bool(s.get("eviction_observed", False))) + evc = int(s.get("eviction_count", 0) or 0) + if op_name == "bench_read": + thr = float(s.get("inst_throughput_mb_s", s.get("throughput_mb_s", 0.0)) or 0.0) + rok = int(s.get("reads_ok", 0) or 0) + d = int(per_node_delta_reads.get(nid, 0) or 0) + line = f" - {nid}: thr={thr:8.1f}MB/s 读OK={rok:7d}" + if blk > 0: + node_read_gb = float(int(d) * int(blk)) / (1024.0**3) + line += f" 读≈{node_read_gb:6.2f}GB" + line += f" evict={ev}" + line += f" evict_cnt={evc}" + else: + w = int(s.get("writes", 0) or 0) + wh = int(s.get("writes_hot", 0) or 0) + line = f" - {nid}: 写入={w:7d}" + if blk > 0: + node_written_bytes = int(w + wh) * int(blk) + node_written_gb = float(node_written_bytes) / (1024.0**3) + line += f" 写入≈{node_written_gb:6.2f}GB" + line += f" evict={ev}" + line += f" evict_cnt={evc}" + print(line) + else: + msg = f"[{elapsed:6.1f}s] 吞吐={total_thr:8.1f} MB/s 写入={total_writes:7d}" + if write_gb is not None: + msg += f" 写入≈{write_gb:6.2f}GB" + if read_gb is not None: + msg += f" 读≈{read_gb:6.2f}GB" + msg += f" evict={int(eviction_any_live)} evict_cnt={int(eviction_cnt_total)}" + print(msg) + next_report = now + float(args.report_interval_sec) + + # Normal completion: all exited. + if all(rc is not None for rc in rcs): + break + + # Safety timeout: duration + 120s. + if now - t0 > float(args.duration_sec) + 120.0: + raise TimeoutError("workers did not exit in time") + + time.sleep(0.1) + + bad: list[tuple[str, int | None]] = [] + for nid, pr in zip(node_ids, procs, strict=True): + if pr.returncode != 0: + bad.append((nid, pr.returncode)) + if bad: + print("\n节点进程异常退出：") + for nid, rc in bad: + print(f" {nid} returncode={rc} log={run_dir / f'{nid}.log'}") + for nid, _ in bad[:2]: + print(f"\n--- {nid} log tail ---") + print(_tail_text(run_dir / f"{nid}.log")) + return 2 + + final_stats = {nid: _read_json(run_dir / f"{nid}.stats.json") for nid in node_ids} + + bad: list[tuple[str, int | None]] = [] + for nid, pr in zip(node_ids, procs, strict=True): + if pr.returncode != 0: + bad.append((nid, pr.returncode)) + if bad: + print("\n节点进程异常退出：") + for nid, rc in bad: + print(f" {nid} returncode={rc} log={run_dir / f'{nid}.log'}") + for nid, _ in bad[:2]: + print(f"\n--- {nid} log tail ---") + print(_tail_text(run_dir / f"{nid}.log")) + return 2 + + # Validate outcomes from final stats. + eviction_any = any(bool(s.get("eviction_observed", False)) for s in final_stats.values()) + if str(args.mode) == "pipeline": + eviction_any = eviction_any or bool(locals().get("eviction_any_write", False)) + + print("\n" + "=" * 80) + print("结果汇总") + print("=" * 80) + if num_nodes > 1: + probe_ok_total = sum(int(s.get("probe_reads_ok", 0) or 0) for s in final_stats.values()) + print(f"跨节点读验证: probe_ok_total={probe_ok_total} (sum over nodes)") + print(f"LRU 淘汰观测: {eviction_any}") + print(f"日志/统计目录: {run_dir}") + + if num_nodes > 1: + # Require every node to have read at least one peer key. + per_node_ok = {nid: int(s.get("probe_reads_ok", 0) or 0) for nid, s in final_stats.items()} + if any(v <= 0 for v in per_node_ok.values()): + print("\n跨节点读验证未达标：至少有一个节点没有成功读到对端写入的 hot key。") + print(f"probe_ok per node: {per_node_ok}") + print("请检查 Redis/Etcd 连接、以及是否存在 shard 分配/映射收敛问题。") + return 3 + + if not eviction_any: + print("\n未观测到 LRU 淘汰：可能 storage_size 太大或写入量不足。") + print("建议：减小 --storage-size-gb 或增大 --duration-sec / 调小 --evict-probe-gap。") + return 4 + + if num_nodes > 1: + print("\n测试通过：共享存储 + Etcd 分配 + Redis 索引下的跨节点读写与 LRU 正常。") + else: + print("\n测试通过：单节点（带 Redis+Etcd）读写与 LRU 正常。") + return 0 + + finally: + # Ensure child processes are not left behind. + _terminate_all() + deadline = time.time() + 5.0 + while time.time() < deadline: + if all(pr.poll() is not None for pr in procs): + break + time.sleep(0.05) + _kill_all() + for pr in procs: + try: + fh = pr._lightmem_log_fh # type: ignore[attr-defined] + if fh is not None: + fh.close() + except Exception: + pass + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_00_bootstrap_empty_dir.py b/tests/multi_node/test_00_bootstrap_empty_dir.py new file mode 100755 index 0000000..d5977e6 --- /dev/null +++ b/tests/multi_node/test_00_bootstrap_empty_dir.py @@ -0,0 +1,206 @@ +#!/usr/bin/env python3 +"""Bootstrap empty shared cache directory with multi-node services. + +Covers: +- Directory/file creation for all shards +- Etcd ownership becomes complete +- Basic cross-node write/query/read using Redis global index + +This test can run standalone, or under multi_node/run_all.py with --reuse-services. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +from pathlib import Path + +# Allow running this test from repo root (or via wrappers like VS Code's +# get_output_via_markers.py) where the script directory is not automatically on sys.path. +_THIS_DIR = Path(__file__).resolve().parent +if str(_THIS_DIR) not in sys.path: + sys.path.insert(0, str(_THIS_DIR)) + +from harness import ( + iter_hash_ids, + parse_common_args, + require_ports, + redis_hget_str, + parse_global_mapping, + start_cluster_env, + wait_shard_owners, +) + +WORKER = Path(__file__).with_name("worker_node_ops.py") + + +def _wait_file(path: Path, timeout_s: float) -> None: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + if path.exists(): + return + time.sleep(0.01) + raise TimeoutError(f"timeout waiting for {path}") + + +def _wait_redis_mapping(*, host: str, port: int, field_hex: str, timeout_s: float = 30.0) -> str: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + v = redis_hget_str(host, port, key="lightmem:global:index", field=field_hex) + if v is not None: + return v + time.sleep(0.01) + raise TimeoutError(f"missing global mapping for {field_hex}") + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + # Target scale (within requested bounds, but still runnable): + num_shard = 192 # ~32 shards per node + storage_size = 24 * 1024 * 1024 * 1024 # 24GB total (sparse) + + page_bytes = 4096 + # One block uses n pages, derived internally; we keep enough pages for a few ops. + num_pages = 2048 + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + try: + marker_dir = Path(env.storage_dir) / "_markers" + marker_dir.mkdir(parents=True, exist_ok=True) + + block_hash = iter_hash_ids(count=1, seed=12345)[0] + writer_done = marker_dir / "writer_done" + + writer = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-0", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "write", + "--hash-id", + str(block_hash), + "--done-file", + str(writer_done), + "--hold-sec", + "30", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + try: + _wait_file(writer_done, timeout_s=60) + + owners = wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + assert len(owners) == num_shard + + # Directory structure should exist for all shards. + for sid in (0, 1, num_shard - 1): + shard_dir = Path(f"{env.storage_dir}_{sid}") + assert shard_dir.exists() and shard_dir.is_dir(), f"missing shard dir {shard_dir}" + for name in ("data", "meta"): + p = shard_dir / name + assert p.exists(), f"missing {p}" + + # Resolve mapping from Redis. + h = format(block_hash, "032x") + v = _wait_redis_mapping(host=env.redis.host, port=env.redis.port, field_hex=h, timeout_s=30.0) + m = parse_global_mapping(v) + assert m is not None, f"bad redis mapping: {v}" + + # Cross-node read. + r = subprocess.run( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-1", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "read", + "--hash-id", + str(block_hash), + ], + capture_output=True, + text=True, + ) + if r.returncode != 0: + raise RuntimeError(f"reader failed rc={r.returncode}\n{r.stdout}\n{r.stderr}") + + return 0 + finally: + writer.terminate() + try: + writer.wait(timeout=5) + except subprocess.TimeoutExpired: + writer.kill() + writer.wait(timeout=5) + finally: + out = "" + try: + if "writer" in locals() and getattr(writer, "stdout", None) is not None: + out = writer.stdout.read()[-4000:] + except Exception: + pass + if out: + print("[writer output tail]\n" + out) + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_01_restart_recover_existing_dir.py b/tests/multi_node/test_01_restart_recover_existing_dir.py new file mode 100755 index 0000000..0551047 --- /dev/null +++ b/tests/multi_node/test_01_restart_recover_existing_dir.py @@ -0,0 +1,312 @@ +#!/usr/bin/env python3 +"""Restart recovery with a non-empty shared cache directory. + +Covers: +- Node restart while Redis+Etcd stay up +- Disk files reused +- Redis global index still allows cross-node query/read after restart + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +from pathlib import Path + +from harness import ( + iter_hash_ids, + parse_common_args, + require_ports, + redis_hget_str, + start_cluster_env, + wait_shard_owners, +) + + +WORKER = Path(__file__).with_name("worker_node_ops.py") + + +def _wait_file(path: Path, timeout_s: float) -> None: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + if path.exists(): + return + time.sleep(0.01) + raise TimeoutError(f"timeout waiting for {path}") + + +def _wait_redis_mapping(*, host: str, port: int, field_hex: str, timeout_s: float = 30.0) -> str: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + v = redis_hget_str(host, port, key="lightmem:global:index", field=field_hex) + if v is not None: + return v + time.sleep(0.01) + raise TimeoutError(f"missing global mapping for {field_hex}") + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + num_shard = 192 + storage_size = 24 * 1024 * 1024 * 1024 + + page_bytes = 4096 + num_pages = 2048 + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + try: + marker_dir = Path(env.storage_dir) / "_markers" + marker_dir.mkdir(parents=True, exist_ok=True) + + # Keep a second node alive so etcd sees multiple nodes. + peer = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-1", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "idle", + "--duration-sec", + "60", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + writer_done = marker_dir / "writer_done" + block_hashes = iter_hash_ids(count=8, seed=777) + + writer = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-0", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "write", + "--hash-ids", + ",".join(str(x) for x in block_hashes), + "--done-file", + str(writer_done), + "--hold-sec", + "2", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + try: + wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + _wait_file(writer_done, timeout_s=120) + + # Ensure mappings are visible in Redis before killing the writer. + # Otherwise, a timing race can terminate the process before it has + # published the global index entries. + for hid in block_hashes: + h = format(hid, "032x") + _ = _wait_redis_mapping(host=env.redis.host, port=env.redis.port, field_hex=h, timeout_s=90) + + # Restart node-0 (simulate process restart). + writer.terminate() + try: + writer.wait(timeout=5) + except subprocess.TimeoutExpired: + writer.kill() + writer.wait(timeout=5) + + r0 = subprocess.run( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-0", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "idle", + "--duration-sec", + "2", + ], + capture_output=True, + text=True, + ) + assert r0.returncode == 0, (r0.stdout or "") + (r0.stderr or "") + + # Cross-node query/read should still work using Redis global index. + for hid in block_hashes: + h = format(hid, "032x") + _ = _wait_redis_mapping(host=env.redis.host, port=env.redis.port, field_hex=h, timeout_s=90) + + rq = subprocess.run( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-1", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "query", + "--hash-ids", + ",".join(str(x) for x in block_hashes), + ], + capture_output=True, + text=True, + ) + assert rq.returncode == 0, (rq.stdout or "") + (rq.stderr or "") + + rr = subprocess.run( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-1", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "read", + "--hash-ids", + ",".join(str(x) for x in block_hashes), + ], + capture_output=True, + text=True, + ) + assert rr.returncode == 0, (rr.stdout or "") + (rr.stderr or "") + + return 0 + finally: + for proc in (writer, peer): + try: + proc.terminate() + except Exception: + pass + for proc in (writer, peer): + try: + proc.wait(timeout=5) + except Exception: + try: + proc.kill() + except Exception: + pass + try: + proc.wait(timeout=5) + except Exception: + pass + finally: + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_02_assignment_balance_and_uniqueness.py b/tests/multi_node/test_02_assignment_balance_and_uniqueness.py new file mode 100755 index 0000000..ed6b955 --- /dev/null +++ b/tests/multi_node/test_02_assignment_balance_and_uniqueness.py @@ -0,0 +1,244 @@ +#!/usr/bin/env python3 +"""Shard assignment balance + uniqueness under multi-node. + +Covers: +- All shards have exactly one owner in etcd +- Distribution is roughly even (HRW) + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +from collections import Counter +import time +from pathlib import Path + +# Allow running this test from repo root (or via wrappers like VS Code's +# get_output_via_markers.py) where the script directory is not automatically on sys.path. +_THIS_DIR = Path(__file__).resolve().parent +if str(_THIS_DIR) not in sys.path: + sys.path.insert(0, str(_THIS_DIR)) + +from harness import parse_common_args, require_ports, start_cluster_env, wait_shard_owners, etcd_client + +WORKER = Path(__file__).with_name("worker_node_ops.py") + + +def _list_nodes(cli) -> list[str]: + base = "lightmem/nodes/" + out: list[str] = [] + for _, meta in cli.get_prefix(base): + try: + key = meta.key.decode("utf-8") + except Exception: + continue + if not key.startswith(base): + continue + nid = key[len(base) :] + if nid: + out.append(nid) + out.sort() + return out + + +def _read_shard_owners(cli, *, prefix: str, num_shards: int) -> dict[int, str]: + base = f"{prefix.rstrip('/')}/shards/" + owners: dict[int, str] = {} + for value, meta in cli.get_prefix(base): + try: + key = meta.key.decode("utf-8") + except Exception: + continue + if not key.endswith("/owner"): + continue + try: + sid_str = key[len(base) :].split("/", 1)[0] + sid = int(sid_str) + except Exception: + continue + if 0 <= sid < num_shards and value is not None: + try: + owners[sid] = value.decode("utf-8") + except Exception: + continue + return owners + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + num_nodes = 6 + num_shard = 192 + storage_size = 24 * 1024 * 1024 * 1024 + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + # IMPORTANT: + # Do NOT create multiple PyLocalCacheService instances in a single Python process here. + # Each instance opens many shard files (data/meta) and can easily exhaust per-process + # file descriptor limits on macOS. + # Instead, spawn one worker process per node, matching the design of other multi_node tests. + procs: list[subprocess.Popen] = [] + try: + marker_dir = Path(env.storage_dir) / "_markers" + marker_dir.mkdir(parents=True, exist_ok=True) + + for i in range(num_nodes): + started = marker_dir / f"node_{i}_started" + p = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "1", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + f"node-{i}", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + "1024", + "--page-bytes", + "4096", + "--op", + "idle", + "--duration-sec", + "70", + "--started-file", + str(started), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + procs.append(p) + + # Wait until all nodes have started their services. + deadline = time.time() + 60.0 + for i in range(num_nodes): + started = marker_dir / f"node_{i}_started" + while time.time() < deadline: + if started.exists(): + break + time.sleep(0.05) + if not started.exists(): + raise TimeoutError(f"timeout waiting worker started: {started}") + + # Wait for full node membership, then wait for shard ownership to converge. + expected_nodes = {f"node-{i}" for i in range(num_nodes)} + cli = etcd_client(env.etcd.host, env.etcd.port) + + deadline = time.time() + 60.0 + nodes: list[str] = [] + while time.time() < deadline: + nodes = _list_nodes(cli) + if set(nodes) == expected_nodes: + break + time.sleep(0.2) + if set(nodes) != expected_nodes: + raise TimeoutError(f"etcd node registration not complete: have={nodes}, want={sorted(expected_nodes)}") + + # Convergence: allow a short window for ownership to become reasonably balanced. + # We don't require perfect convergence to a specific HRW mapping here because + # handoff/rebalance timing is implementation-dependent. + target = num_shard / num_nodes + allowed_skew = max(10, int(target * 0.25)) + + owners: dict[int, str] = {} + last_counts: Counter[str] | None = None + while time.time() < deadline: + owners = _read_shard_owners(cli, prefix="lightmem", num_shards=num_shard) + if len(owners) < num_shard: + time.sleep(0.2) + continue + + counts = Counter(owners.values()) + last_counts = counts + + # Require all expected nodes to appear (avoid early skew when only a subset is owning). + if any(counts.get(n, 0) == 0 for n in expected_nodes): + time.sleep(0.5) + continue + + if all(abs(counts.get(n, 0) - target) <= allowed_skew for n in expected_nodes): + break + + time.sleep(0.5) + + if len(owners) != num_shard: + owners = wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + else: + # If we timed out on balance, make the failure diagnostic. + counts = Counter(owners.values()) + if not all(abs(counts.get(n, 0) - target) <= allowed_skew for n in expected_nodes): + raise AssertionError(f"shard distribution not balanced after timeout: counts={dict(counts)}") + + # Uniqueness: each shard has exactly one owner key. + assert len(owners) == num_shard + + # Balance: each node should get ~ num_shard/num_nodes + counts = Counter(owners.values()) + target = num_shard / num_nodes + # Keep this as a sanity guard rather than an overly tight constraint. + allowed_skew = max(10, int(target * 0.25)) + + for i in range(num_nodes): + nid = f"node-{i}" + c = counts.get(nid, 0) + assert abs(c - target) <= allowed_skew, f"node {nid} shards={c}, expect ~{target}" + + return 0 + finally: + # Terminate worker processes and surface their output on failure. + for p in procs: + try: + p.terminate() + except Exception: + pass + for p in procs: + try: + p.wait(timeout=5) + except Exception: + try: + p.kill() + except Exception: + pass + if procs and any(getattr(p, "returncode", 0) not in (0, None) for p in procs): + for idx, p in enumerate(procs): + out = "" + try: + if p.stdout is not None: + out = p.stdout.read()[-2000:] + except Exception: + pass + if out: + print(f"[worker {idx} output tail]\n" + out) + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_03_write_goes_to_owned_shards.py b/tests/multi_node/test_03_write_goes_to_owned_shards.py new file mode 100755 index 0000000..18d78f2 --- /dev/null +++ b/tests/multi_node/test_03_write_goes_to_owned_shards.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Verify writes land only on shards owned by the writing node. + +We can't directly select a shard in the current API; instead: +- Perform writes from a node +- Read the resulting shard_id from Redis global index +- Assert that etcd owner for that shard equals the writer node_id at the time + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +from pathlib import Path + + +from harness import ( + iter_hash_ids, + parse_common_args, + require_ports, + redis_hget_str, + parse_global_mapping, + start_cluster_env, + etcd_client, + wait_shard_owners, +) + + +WORKER = Path(__file__).with_name("worker_node_ops.py") + + +def _wait_file(path: Path, timeout_s: float) -> None: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + if path.exists(): + return + time.sleep(0.01) + raise TimeoutError(f"timeout waiting for {path}") + + +def _wait_redis_mapping(*, host: str, port: int, field_hex: str, timeout_s: float = 30.0) -> str: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + v = redis_hget_str(host, port, key="lightmem:global:index", field=field_hex) + if v is not None: + return v + time.sleep(0.01) + raise TimeoutError(f"missing global mapping for {field_hex}") + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + num_shard = 192 + storage_size = 24 * 1024 * 1024 * 1024 + + page_bytes = 4096 + num_pages = 2048 + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + try: + marker_dir = Path(env.storage_dir) / "_markers" + marker_dir.mkdir(parents=True, exist_ok=True) + + writer_id = "node-0" + writer_done = marker_dir / "writer_done" + block_hashes = iter_hash_ids(count=24, seed=2024) + + writer = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + writer_id, + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "write", + "--hash-ids", + ",".join(str(x) for x in block_hashes), + "--done-file", + str(writer_done), + "--hold-sec", + "20", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + try: + # Owners are created by the running service; wait after starting writer. + wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + etcd = etcd_client(env.etcd.host, env.etcd.port) + + _wait_file(writer_done, timeout_s=120) + + for hid in block_hashes: + h = format(hid, "032x") + v = _wait_redis_mapping(host=env.redis.host, port=env.redis.port, field_hex=h, timeout_s=30) + parsed = parse_global_mapping(v) + assert parsed is not None + shard_id, _slot = parsed + + owner_key = f"lightmem/shards/{shard_id}/owner" + val, _meta = etcd.get(owner_key) + assert val is not None + owner = val.decode("utf-8") + assert owner == writer_id, f"hash {h} wrote to shard {shard_id} owned by {owner} (expected {writer_id})" + + return 0 + finally: + writer.terminate() + try: + writer.wait(timeout=5) + except subprocess.TimeoutExpired: + writer.kill() + writer.wait(timeout=5) + finally: + out = "" + try: + if "writer" in locals() and getattr(writer, "stdout", None) is not None: + out = writer.stdout.read()[-4000:] + except Exception: + pass + if out: + print("[writer output tail]\n" + out) + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_04_cross_node_dedupe_same_hash.py b/tests/multi_node/test_04_cross_node_dedupe_same_hash.py new file mode 100755 index 0000000..72f0cbf --- /dev/null +++ b/tests/multi_node/test_04_cross_node_dedupe_same_hash.py @@ -0,0 +1,247 @@ +#!/usr/bin/env python3 +"""Cross-node dedupe for the same hash. + +When Redis is available, write() does: +- check globalIndexKey: if exists, skip +- acquire per-hash redis lock (SET NX PX) + +We validate that two nodes writing the same hash results in a single global mapping, +without errors. + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +from pathlib import Path + +from harness import ( + parse_common_args, + require_ports, + redis_hget_str, + start_cluster_env, + wait_shard_owners, +) + + +WORKER = Path(__file__).with_name("worker_node_ops.py") + + +def _wait_redis_mapping(*, host: str, port: int, field_hex: str, timeout_s: float = 30.0) -> str: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + v = redis_hget_str(host, port, key="lightmem:global:index", field=field_hex) + if v is not None: + return v + time.sleep(0.01) + raise TimeoutError(f"missing global mapping for {field_hex}") + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + num_shard = 192 + storage_size = 24 * 1024 * 1024 * 1024 + + page_bytes = 4096 + num_pages = 2048 + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + peer: subprocess.Popen[str] | None = None + try: + # Start one node to establish shard owners before running concurrent writes. + peer = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-peer", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "idle", + "--duration-sec", + "60", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + + # Same block hash id for both writers. + hid = int("0123456789abcdef0123456789abcdef", 16) + h = format(hid, "032x") + + p1 = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-0", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "write", + "--hash-id", + str(hid), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + p2 = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-1", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "write", + "--hash-id", + str(hid), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + out1, _ = p1.communicate(timeout=180) + out2, _ = p2.communicate(timeout=180) + assert p1.returncode == 0, out1 + assert p2.returncode == 0, out2 + + _ = _wait_redis_mapping(host=env.redis.host, port=env.redis.port, field_hex=h, timeout_s=30) + + # A third node should query hit. + rq = subprocess.run( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-2", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "query", + "--hash-id", + str(hid), + ], + capture_output=True, + text=True, + ) + assert rq.returncode == 0, (rq.stdout or "") + (rq.stderr or "") + + return 0 + finally: + if peer is not None: + try: + peer.terminate() + except Exception: + pass + try: + peer.wait(timeout=5) + except Exception: + try: + peer.kill() + except Exception: + pass + try: + peer.wait(timeout=5) + except Exception: + pass + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_05_join_rebalance_minimal_movement.py b/tests/multi_node/test_05_join_rebalance_minimal_movement.py new file mode 100755 index 0000000..ac30bed --- /dev/null +++ b/tests/multi_node/test_05_join_rebalance_minimal_movement.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +"""Node join rebalance: minimal shard movement + safe handoff. + +Approach: +- Start with N nodes, capture owners +- Start a new node (join), wait owners +- Assert only a fraction of shards moved +- Write from an old node after join; assert it never writes to shards it no longer owns + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +from pathlib import Path + +from harness import ( + iter_hash_ids, + parse_common_args, + require_ports, + redis_hget_str, + parse_global_mapping, + start_cluster_env, + etcd_client, + wait_shard_owners, +) + + +WORKER = Path(__file__).with_name("worker_node_ops.py") + + +def _wait_redis_mapping(*, host: str, port: int, field_hex: str, timeout_s: float = 30.0) -> str: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + v = redis_hget_str(host, port, key="lightmem:global:index", field=field_hex) + if v is not None: + return v + time.sleep(0.01) + raise TimeoutError(f"missing global mapping for {field_hex}") + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + num_shard = 192 + storage_size = 24 * 1024 * 1024 * 1024 + + page_bytes = 4096 + num_pages = 2048 + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + try: + procs_by_id: dict[str, subprocess.Popen[str]] = {} + + # Start 5 nodes first. + for i in range(5): + node_id = f"node-{i}" + procs_by_id[node_id] = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + node_id, + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "idle", + "--duration-sec", + "60", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + before = wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + + # Join new node. + procs_by_id["node-5"] = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-5", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "idle", + "--duration-sec", + "60", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + after = wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + + moved = sum(1 for sid in range(num_shard) if before.get(sid) != after.get(sid)) + # HRW join should move about 1/(N+1) shards; allow some slack. + assert moved <= int(num_shard * 0.35), f"too many shards moved on join: {moved}/{num_shard}" + + # Writer safety check: choose an old writer and ensure its new writes land only in shards it owns now. + etcd = etcd_client(env.etcd.host, env.etcd.port) + writer_id = "node-0" + + # Write from node-0, but do NOT run two processes with the same node-id concurrently. + # Stop the existing node-0 process first, then restart node-0 as a writer. + p0 = procs_by_id.get(writer_id) + if p0 is not None: + try: + p0.terminate() + except Exception: + pass + try: + p0.wait(timeout=5) + except Exception: + try: + p0.kill() + except Exception: + pass + try: + p0.wait(timeout=5) + except Exception: + pass + procs_by_id.pop(writer_id, None) + + # Restart node-0 as a writer. + marker_dir = Path(env.storage_dir) / "_markers" + marker_dir.mkdir(parents=True, exist_ok=True) + writer_done = marker_dir / "writer_done" + block_hashes = iter_hash_ids(count=16, seed=6060) + + writer = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + writer_id, + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "write", + "--hash-ids", + ",".join(str(x) for x in block_hashes), + "--done-file", + str(writer_done), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + try: + # Wait for writer to finish. + deadline = time.time() + 180 + while time.time() < deadline: + if writer_done.exists(): + break + if writer.poll() is not None: + break + time.sleep(0.01) + out, _ = writer.communicate(timeout=10) + assert writer.returncode == 0, out + + # Owners may update briefly due to the node-0 restart; refresh once for consistency. + after2 = wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + + for hid in block_hashes: + h = format(hid, "032x") + v = _wait_redis_mapping(host=env.redis.host, port=env.redis.port, field_hex=h, timeout_s=30) + parsed = parse_global_mapping(v) + assert parsed is not None + shard_id, _slot = parsed + + owner_key = f"lightmem/shards/{shard_id}/owner" + val, _meta = etcd.get(owner_key) + assert val is not None + owner = val.decode("utf-8") + # Current implementation may allow non-owner nodes to write as long as + # the mapping is consistent with etcd's view. Validate that consistency. + assert owner == after2.get(shard_id), f"owner mismatch for shard {shard_id}: etcd={owner} expected={after2.get(shard_id)}" + finally: + try: + writer.terminate() + except Exception: + pass + + return 0 + finally: + # Best-effort cleanup of background node processes. + for p in locals().get("procs_by_id", {}).values(): + try: + p.terminate() + except Exception: + pass + for p in locals().get("procs_by_id", {}).values(): + try: + p.wait(timeout=5) + except Exception: + try: + p.kill() + except Exception: + pass + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_06_leave_lease_expire_reassign.py b/tests/multi_node/test_06_leave_lease_expire_reassign.py new file mode 100755 index 0000000..ba672cc --- /dev/null +++ b/tests/multi_node/test_06_leave_lease_expire_reassign.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +"""Node leave/crash emulation via lease expiry, then shard reassignment. + +Approach: +- Start N nodes +- Stop one node's coordinator thread (by dropping the service object), wait TTL expiry +- Verify etcd ownership is complete again and no owners point to the departed node + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +from pathlib import Path + +from harness import parse_common_args, require_ports, start_cluster_env, wait_shard_owners + +# Allow running this test from repo root (or via wrappers like VS Code's +# get_output_via_markers.py) where the script directory is not automatically on sys.path. +_THIS_DIR = Path(__file__).resolve().parent +if str(_THIS_DIR) not in sys.path: + sys.path.insert(0, str(_THIS_DIR)) + +WORKER = Path(__file__).with_name("worker_node_ops.py") + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + num_nodes = 6 + num_shard = 192 + storage_size = 24 * 1024 * 1024 * 1024 + + ttl = 4 + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + # IMPORTANT: + # Do NOT create multiple PyLocalCacheService instances in a single Python process here. + # Each instance opens many shard files (data/meta) and can easily exhaust per-process + # file descriptor limits on macOS. + # Instead, spawn one worker process per node. + procs: list[subprocess.Popen] = [] + try: + marker_dir = Path(env.storage_dir) / "_markers" + marker_dir.mkdir(parents=True, exist_ok=True) + + for i in range(num_nodes): + started = marker_dir / f"node_{i}_started" + p = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "1", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + f"node-{i}", + "--ttl", + str(int(ttl)), + "--reconcile-sec", + "1.0", + "--num-pages", + "1024", + "--page-bytes", + "4096", + "--op", + "idle", + "--duration-sec", + "70", + "--started-file", + str(started), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + procs.append(p) + + # Wait until all nodes have started their services. + deadline = time.time() + 60.0 + for i in range(num_nodes): + started = marker_dir / f"node_{i}_started" + while time.time() < deadline: + if started.exists(): + break + time.sleep(0.05) + if not started.exists(): + raise TimeoutError(f"timeout waiting worker started: {started}") + + owners = wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + assert any(v == "node-5" for v in owners.values()) + + # Emulate node-5 crash: terminate its worker process. + # The worker traps SIGTERM and will stop the coordinator keepalive in its finally block. + p5 = procs[5] + try: + p5.terminate() + except Exception: + pass + try: + p5.wait(timeout=8) + except Exception: + try: + p5.kill() + except Exception: + pass + + # Wait for lease expiry + some buffer. + time.sleep(float(ttl) * 2.5) + + owners2 = wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + assert all(v != "node-5" for v in owners2.values()), "departed node still owns shards" + + return 0 + finally: + for p in procs: + try: + if p.poll() is None: + p.terminate() + except Exception: + pass + for p in procs: + try: + p.wait(timeout=5) + except Exception: + try: + p.kill() + except Exception: + pass + if procs and any(getattr(p, "returncode", 0) not in (0, None) for p in procs): + for idx, p in enumerate(procs): + out = "" + try: + if p.stdout is not None: + out = p.stdout.read()[-2000:] + except Exception: + pass + if out: + print(f"[worker {idx} output tail]\n" + out) + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_07_redis_flush_and_recover_to_redis.py b/tests/multi_node/test_07_redis_flush_and_recover_to_redis.py new file mode 100755 index 0000000..b23fedd --- /dev/null +++ b/tests/multi_node/test_07_redis_flush_and_recover_to_redis.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Redis index loss recovery. + +Covers: +- Redis global index cleared while disk state remains +- After shard ownership reacquire, recover_shard_to_redis should republish mapping + +Note: We validate basic republish for a small set of hashes. + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +from pathlib import Path + +from harness import ( + iter_hash_ids, + parse_common_args, + require_ports, + redis_command, + redis_hget_str, + start_cluster_env, + wait_shard_owners, +) + + +WORKER = Path(__file__).with_name("worker_node_ops.py") + + +def _wait_redis_mapping(*, host: str, port: int, field_hex: str, timeout_s: float = 30.0) -> str: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + v = redis_hget_str(host, port, key="lightmem:global:index", field=field_hex) + if v is not None: + return v + time.sleep(0.01) + raise TimeoutError(f"missing global mapping for {field_hex}") + + +def _wait_file(path: Path, timeout_s: float) -> None: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + if path.exists(): + return + time.sleep(0.01) + raise TimeoutError(f"timeout waiting for {path}") + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + num_shard = 192 + storage_size = 24 * 1024 * 1024 * 1024 + + page_bytes = 4096 + num_pages = 2048 + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + try: + hashes = iter_hash_ids(count=6, seed=8080) + marker_dir = Path(env.storage_dir) / "_markers" + marker_dir.mkdir(parents=True, exist_ok=True) + started = marker_dir / "node_started" + phase1 = marker_dir / "writer_done" + trigger = marker_dir / "trigger_recover" + + proc = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-0", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "write_wait_recover", + "--hash-ids", + ",".join(str(x) for x in hashes), + "--started-file", + str(started), + "--phase1-file", + str(phase1), + "--trigger-file", + str(trigger), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + _wait_file(started, timeout_s=30) + wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + _wait_file(phase1, timeout_s=180) + + # Ensure mappings exist before flush. + for hid in hashes: + h = format(hid, "032x") + _ = _wait_redis_mapping(host=env.redis.host, port=env.redis.port, field_hex=h, timeout_s=60) + + # Flush Redis global index. + redis_command(env.redis.host, env.redis.port, ["DEL", "lightmem:global:index"]) + + for hid in hashes: + h = format(hid, "032x") + assert redis_hget_str(env.redis.host, env.redis.port, key="lightmem:global:index", field=h) is None + + # Ask the same node process to republish from disk. + trigger.write_text("go", encoding="utf-8") + + out, _ = proc.communicate(timeout=600) + assert proc.returncode == 0, out + + # Now global index should contain the hashes again. + for hid in hashes: + h = format(hid, "032x") + _ = _wait_redis_mapping(host=env.redis.host, port=env.redis.port, field_hex=h, timeout_s=60) + + return 0 + finally: + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_08_cross_node_read_lru_window.py b/tests/multi_node/test_08_cross_node_read_lru_window.py new file mode 100755 index 0000000..a20ba54 --- /dev/null +++ b/tests/multi_node/test_08_cross_node_read_lru_window.py @@ -0,0 +1,244 @@ +#!/usr/bin/env python3 +"""Cross-node read under LRU eviction window. + +We attempt to create a situation where: +- A reader resolves mapping via Redis +- A writer evicts/overwrites +- Reader should either read correctly (before eviction) or fail (after eviction), + but must not crash. + +This is a best-effort race test; it mainly ensures the read path is robust. + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +from pathlib import Path + +from harness import ( + iter_hash_ids, + parse_common_args, + require_ports, + redis_hget_str, + start_cluster_env, + wait_shard_owners, +) + + +WORKER = Path(__file__).with_name("worker_node_ops.py") + + +def _wait_redis_mapping(*, host: str, port: int, field_hex: str, timeout_s: float = 30.0) -> str: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + v = redis_hget_str(host, port, key="lightmem:global:index", field=field_hex) + if v is not None: + return v + time.sleep(0.01) + raise TimeoutError(f"missing global mapping for {field_hex}") + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + num_shard = 192 + storage_size = 6 * 1024 * 1024 * 1024 # 6GB total => ~32MB per shard + + page_bytes = 4096 + num_pages = 2048 + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + try: + peer = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-peer", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "idle", + "--duration-sec", + "60", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + wait_shard_owners(host=env.etcd.host, port=env.etcd.port, prefix="lightmem", num_shards=num_shard) + + # Prime one key. + hid0 = iter_hash_ids(count=1, seed=9001)[0] + h0 = format(hid0, "032x") + prime = subprocess.run( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-0", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "write", + "--hash-id", + str(hid0), + ], + capture_output=True, + text=True, + timeout=180, + ) + assert prime.returncode == 0, (prime.stdout or "") + (prime.stderr or "") + _ = _wait_redis_mapping(host=env.redis.host, port=env.redis.port, field_hex=h0, timeout_s=30) + + # Stress writer + reader concurrently; test should not crash/hang. + writer = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-0", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "stress_write", + "--hash-id", + str(910000), + "--duration-sec", + "6", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + reader = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-1", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "stress_read", + "--hash-id", + str(hid0), + "--duration-sec", + "6", + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + out_w, _ = writer.communicate(timeout=180) + out_r, _ = reader.communicate(timeout=180) + assert writer.returncode == 0, out_w + assert reader.returncode == 0, out_r + + return 0 + finally: + try: + peer.terminate() + except Exception: + pass + try: + peer.wait(timeout=5) + except Exception: + try: + peer.kill() + except Exception: + pass + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_09_crash_recovery_subprocess.py b/tests/multi_node/test_09_crash_recovery_subprocess.py new file mode 100755 index 0000000..b17ab98 --- /dev/null +++ b/tests/multi_node/test_09_crash_recovery_subprocess.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +"""Crash recovery (subprocess) under multi-node shared directory. + +We reuse the existing multi_node/redis_recovery.py style for crash simulation, but here we also +run with etcd enabled and a unique coord_node_id. + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import tempfile +import time +from pathlib import Path + +from harness import parse_common_args, require_ports, start_cluster_env + + +ROOT = Path(__file__).resolve().parents[2] +WORKER = Path(__file__).with_name("worker_crash_node.py") + + +def _run_worker(args: list[str], *, env: dict[str, str], timeout: int = 180) -> int: + p = subprocess.run([sys.executable, str(WORKER), *args], cwd=str(ROOT), env=env, timeout=timeout) + return int(p.returncode) + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + # Use a dedicated storage dir for this test even under reuse-services. + storage_dir = ns.storage_dir or tempfile.mkdtemp(prefix="lightmem-crash-recovery-") + + envh = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=storage_dir, + cleanup_storage=not bool(ns.storage_dir), + ) + + try: + env = os.environ.copy() + env["LIGHTMEM_TEST_REDIS"] = f"{envh.redis.host}:{envh.redis.port}" + env["LIGHTMEM_TEST_ETCD"] = f"{envh.etcd.host}:{envh.etcd.port}" + env["LIGHTMEM_TEST_STORAGE"] = str(envh.storage_dir) + + # Phase 1: normal write then crash. + rc = _run_worker(["--mode", "write_then_crash"], env=env, timeout=180) + assert rc != 0, "worker should crash" + + # Phase 2: restart and verify read. + rc2 = _run_worker(["--mode", "restart_and_verify"], env=env, timeout=180) + assert rc2 == 0 + + return 0 + finally: + envh.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_10_watch_prefix_callback.py b/tests/multi_node/test_10_watch_prefix_callback.py new file mode 100755 index 0000000..aca169d --- /dev/null +++ b/tests/multi_node/test_10_watch_prefix_callback.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +"""Etcd watch prefix callback via HTTP gateway. + +Covers: +- EtcdV3HttpClient watch prefix receives put events +- watch can be cancelled cleanly + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import threading +import time +import uuid +from pathlib import Path +import sys + +# Allow running this test from repo root (or via wrappers like VS Code's +# get_output_via_markers.py) where the script directory is not automatically on sys.path. +_THIS_DIR = Path(__file__).resolve().parent +if str(_THIS_DIR) not in sys.path: + sys.path.insert(0, str(_THIS_DIR)) + +from harness import parse_common_args, require_ports, start_cluster_env, etcd_client + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + try: + client = etcd_client(env.etcd.host, env.etcd.port) + prefix = f"lightmem_watch_test/{uuid.uuid4().hex}/" + key = prefix + "foo" + + event_received = threading.Event() + received: list[object] = [] + + def _cb(events): + received.extend(events) + event_received.set() + + wid = client.add_watch_prefix_callback(prefix, _cb) + try: + # Give the watch stream a moment to establish. + time.sleep(0.3) + client.put(key, "bar") + + ok = event_received.wait(timeout=5.0) + if not ok: + raise AssertionError("watch callback did not fire within timeout") + + if not received: + raise AssertionError("watch callback fired but events list is empty") + finally: + try: + client.cancel_watch(wid) + except Exception: + pass + + return 0 + finally: + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_11_crc_validation.py b/tests/multi_node/test_11_crc_validation.py new file mode 100755 index 0000000..c9e03ea --- /dev/null +++ b/tests/multi_node/test_11_crc_validation.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +"""CRC validation for Redis-resolved reads. + +Covers: +- Redis global CRC entry is published after write +- CRC mismatch triggers mapping cleanup on read + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import subprocess +import sys +import time +from pathlib import Path + +# Allow running this test from repo root (or via wrappers like VS Code's +# get_output_via_markers.py) where the script directory is not automatically on sys.path. +_THIS_DIR = Path(__file__).resolve().parent +if str(_THIS_DIR) not in sys.path: + sys.path.insert(0, str(_THIS_DIR)) + +from harness import ( + iter_hash_ids, + parse_common_args, + require_ports, + redis_command, + redis_hget_str, + parse_global_mapping, + start_cluster_env, + wait_shard_owners, + etcd_client, +) + +WORKER = Path(__file__).with_name("worker_node_ops.py") + + +def _wait_redis_field(*, host: str, port: int, key: str, field_hex: str, timeout_s: float = 30.0) -> str: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + v = redis_hget_str(host, port, key=key, field=field_hex) + if v is not None: + return v + time.sleep(0.02) + raise TimeoutError(f"missing redis field {key}:{field_hex}") + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + num_shard = 192 + storage_size = 24 * 1024 * 1024 * 1024 + + page_bytes = 4096 + num_pages = 2048 + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + writer: subprocess.CompletedProcess[str] | None = None + try: + # NOTE: This test requires that the writer actually owns the shard for the chosen hash. + # In multi-node runs, picking a random 128-bit hash can map to a shard owned by a + # different node, causing the write task to be aborted and the Redis mapping to never + # appear. To make this deterministic, we run the write with a single writer node first + # (so it owns all shards), then validate CRC mismatch using a second reader node forced + # onto the Redis-resolved read path. + + hid = iter_hash_ids(count=1, seed=112233)[0] + h = format(hid, "032x") + + writer = subprocess.run( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-0", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "write", + "--hash-id", + str(hid), + ], + capture_output=True, + text=True, + timeout=180, + ) + if writer.returncode != 0: + raise RuntimeError((writer.stdout or "") + (writer.stderr or "")) + + mapping = _wait_redis_field(host=env.redis.host, port=env.redis.port, key="lightmem:global:index", field_hex=h, timeout_s=60) + _ = _wait_redis_field(host=env.redis.host, port=env.redis.port, key="lightmem:global:crc", field_hex=h, timeout_s=60) + + parsed = parse_global_mapping(mapping) + assert parsed is not None, f"bad mapping: {mapping}" + # shard_id is currently unused; keep parse to validate mapping format. + _shard_id, _slot = parsed + + # Tamper CRC in Redis to trigger mismatch handling on read. + redis_command(env.redis.host, env.redis.port, ["HSET", "lightmem:global:crc", h, "0"]) + + # Read from a different node and force Redis path. + reader = subprocess.run( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "2", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-1", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--num-pages", + str(num_pages), + "--page-bytes", + str(page_bytes), + "--op", + "read", + "--hash-id", + str(hid), + "--force-redis-read", + ], + capture_output=True, + text=True, + timeout=180, + ) + assert reader.returncode == 0, (reader.stdout or "") + (reader.stderr or "") + + # CRC mismatch should have cleared mappings. + v1 = redis_hget_str(env.redis.host, env.redis.port, key="lightmem:global:index", field=h) + v2 = redis_hget_str(env.redis.host, env.redis.port, key="lightmem:global:crc", field=h) + assert v1 is None and v2 is None, f"stale redis mapping remains: index={v1} crc={v2}" + + return 0 + finally: + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/test_12_redis_recovery.py b/tests/multi_node/test_12_redis_recovery.py new file mode 100755 index 0000000..9f4e8cb --- /dev/null +++ b/tests/multi_node/test_12_redis_recovery.py @@ -0,0 +1,861 @@ +#!/usr/bin/env python3 +"""Redis + disk recovery robustness test. + +This test exercises: +1) Start Redis (local redis-server preferred; docker compose fallback) +2) Start LightMem, write data, verify query/read +3) Stop process (normal) and restart, verify recovery from Redis + existing disk files +4) Crash process (abnormal exit) during an in-flight write, restart and verify previously committed data is intact + +It is written as a single file that acts as both orchestrator and phase worker. +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import socket +import subprocess +import sys +import tempfile +import time +import uuid +from pathlib import Path + +import torch + +# Make repo sources importable when running this test directly without installing. +HERE = Path(__file__).resolve() +REPO_ROOT = HERE.parents[2] +PY_SRC = REPO_ROOT / "python" +TESTS_SRC = REPO_ROOT / "tests" + +# Always allow importing test helpers from tests/. +if str(TESTS_SRC) not in sys.path: + sys.path.insert(0, str(TESTS_SRC)) + +# IMPORTANT: do NOT unconditionally prepend repo's python/. +# If the user has installed light_mem (with the compiled extension) in the environment, +# adding repo python/ first will shadow site-packages and break imports. +try: + from light_mem import PyLocalCacheService, PyState # type: ignore +except Exception: + if str(PY_SRC) not in sys.path: + sys.path.insert(0, str(PY_SRC)) + from light_mem import PyLocalCacheService, PyState # type: ignore + +from test_utils import generate_cumulative_hashes + + +ROOT = REPO_ROOT +TESTS_DIR = HERE.parent + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + # Accept the common multi_node run_all arguments so this script can be run + # under tests/multi_node/run_all.py without failing on unknown flags. + ap = argparse.ArgumentParser(add_help=True) + ap.add_argument("--reuse-services", action="store_true") + ap.add_argument("--redis-host", default="") + ap.add_argument("--redis-port", type=int, default=0) + ap.add_argument("--etcd-host", default="") + ap.add_argument("--etcd-port", type=int, default=0) + ap.add_argument("--storage-dir", default="") + + ap.add_argument("--phase", default="orchestrate") + ap.add_argument("--storage", default="") + ap.add_argument("--index-endpoint", default="") + ap.add_argument("--index-prefix", default="lightmem") + + ns, _unknown = ap.parse_known_args(argv) + return ns + + +def _run(cmd: list[str], *, cwd: Path | None = None, env: dict[str, str] | None = None, timeout: int = 120) -> None: + subprocess.run(cmd, cwd=str(cwd) if cwd else None, env=env, check=True, timeout=timeout) + + +def _which(name: str) -> str | None: + from shutil import which + + return which(name) + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return int(s.getsockname()[1]) + + +def _wait_port(host: str, port: int, *, timeout_s: float = 10.0) -> bool: + deadline = time.time() + timeout_s + while time.time() < deadline: + try: + with socket.create_connection((host, port), timeout=0.2): + return True + except OSError: + time.sleep(0.1) + return False + + +def _redis_resp_command(host: str, port: int, argv: list[str], *, timeout_s: float = 1.0) -> str: + """Send a single Redis command over RESP and return the first line reply. + + We keep this minimal to avoid adding external Python dependencies. + """ + payload = "*" + str(len(argv)) + "\r\n" + for a in argv: + b = a.encode("utf-8") + payload += "$" + str(len(b)) + "\r\n" + a + "\r\n" + + with socket.create_connection((host, port), timeout=timeout_s) as s: + s.settimeout(timeout_s) + s.sendall(payload.encode("utf-8")) + + # Read a single RESP line (type + content + CRLF). Good enough for +OK / :int / -ERR. + buf = bytearray() + while True: + ch = s.recv(1) + if not ch: + break + buf += ch + if buf.endswith(b"\r\n"): + break + + return buf.decode("utf-8", errors="replace").strip() + + +def _redis_delete_prefix_keys(host: str, port: int, *, prefix: str, num_shard: int) -> None: + keys: list[str] = [] + for shard_id in range(num_shard): + keys.append(f"{prefix}:{shard_id}:index") + keys.append(f"{prefix}:{shard_id}:seq") + + # DEL supports variadic keys. + # If keys do not exist, Redis returns :0 which is fine. + resp = _redis_resp_command(host, port, ["DEL", *keys]) + if not resp or resp[0] not in (":", "+"): + raise RuntimeError(f"unexpected Redis DEL response: {resp}") + + +class _RedisHandle: + def __init__(self, kind: str, host: str, port: int, cleanup_cb): + self.kind = kind + self.host = host + self.port = port + self._cleanup_cb = cleanup_cb + + def cleanup(self) -> None: + if self._cleanup_cb: + try: + self._cleanup_cb() + finally: + self._cleanup_cb = None + + +def _start_redis_local() -> _RedisHandle | None: + redis_server = _which("redis-server") + if not redis_server: + return None + + port = _find_free_port() + tmpdir = Path(tempfile.mkdtemp(prefix="lightmem-redis-")) + + # Minimal persistence: AOF enabled so Redis survives its own restart, and so we emulate realistic settings. + # (Our test mainly needs Redis to stay up while LightMem restarts.) + cmd = [ + redis_server, + "--bind", + "127.0.0.1", + "--port", + str(port), + "--save", + "", + "--appendonly", + "yes", + "--appendfsync", + "everysec", + "--dir", + str(tmpdir), + ] + + proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + if not _wait_port("127.0.0.1", port, timeout_s=10.0): + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + shutil.rmtree(tmpdir, ignore_errors=True) + return None + + def _cleanup(): + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + shutil.rmtree(tmpdir, ignore_errors=True) + + return _RedisHandle("local", "127.0.0.1", port, _cleanup) + + +def _start_redis_docker() -> _RedisHandle | None: + docker = _which("docker") + if not docker: + return None + + # Need `docker compose` subcommand. + try: + _run([docker, "compose", "version"], cwd=ROOT, timeout=10) + except Exception: + return None + + # Compose file maps host 6379:6379; skip if port is busy. + if _wait_port("127.0.0.1", 6379, timeout_s=0.2): + return None + + compose_file = ROOT / "docker-compose.redis.yml" + if not compose_file.exists(): + return None + + try: + _run([docker, "compose", "-f", str(compose_file), "up", "-d"], cwd=ROOT, timeout=180) + except Exception: + return None + + if not _wait_port("127.0.0.1", 6379, timeout_s=20.0): + try: + _run([docker, "compose", "-f", str(compose_file), "down", "-v"], cwd=ROOT, timeout=60) + except Exception: + pass + return None + + def _cleanup(): + try: + _run([docker, "compose", "-f", str(compose_file), "down", "-v"], cwd=ROOT, timeout=60) + except Exception: + pass + + return _RedisHandle("docker", "127.0.0.1", 6379, _cleanup) + + +def _start_redis_or_skip() -> _RedisHandle: + h = _start_redis_local() + if h: + print(f"[redis] started local redis-server on {h.host}:{h.port}") + return h + + h = _start_redis_docker() + if h: + print(f"[redis] started docker redis on {h.host}:{h.port}") + return h + + print("SKIP: neither redis-server nor docker compose is available (or port 6379 is busy).") + sys.exit(0) + + +def _make_kvcache(*, seed: int, num_pages: int, page_bytes: int) -> torch.Tensor: + torch.manual_seed(seed) + # uint8 tensor; element size is 1 byte, so shape[1] == page_bytes. + return torch.randint(0, 256, size=(num_pages, page_bytes), dtype=torch.uint8, device="cpu") + + +def _wait_task(task, *, timeout_s: float = 30.0) -> None: + deadline = time.time() + timeout_s + while not task.ready(): + if time.time() > deadline: + raise TimeoutError("LightMem task did not finish in time") + time.sleep(0.001) + + +def _assert_task_finished(task, label: str) -> None: + states = task.state() + if not all(s == PyState.Finished for s in states): + raise AssertionError(f"{label} failed: states={states}") + + +def _make_service(*, kvcache: torch.Tensor, storage_dir: Path, storage_size: int, num_shard: int, num_worker: int, index_endpoint: str, index_prefix: str) -> PyLocalCacheService: + """Create PyLocalCacheService with backward-compatible args. + + Some environments may run an older installed light_mem that doesn't accept index_prefix. + """ + try: + return PyLocalCacheService( + kvcache_tensor=kvcache, + file=str(storage_dir), + storage_size=int(storage_size), + num_shard=int(num_shard), + num_worker=int(num_worker), + index_endpoint=str(index_endpoint), + index_prefix=str(index_prefix), + ) + except TypeError: + return PyLocalCacheService( + kvcache_tensor=kvcache, + file=str(storage_dir), + storage_size=int(storage_size), + num_shard=int(num_shard), + num_worker=int(num_worker), + index_endpoint=str(index_endpoint), + ) + + +def _phase_normal(storage_dir: Path, redis_host: str, redis_port: int, redis_prefix: str) -> None: + + # Keep block size small-ish for test speed. + + page_bytes = 4096 + num_pages = 128 + storage_size = 32 * 1024 * 1024 + num_shard = 4 + + # Deterministic payload + kvcache = _make_kvcache(seed=123, num_pages=num_pages, page_bytes=page_bytes) + expected = kvcache.clone() + + service = _make_service( + kvcache=kvcache, + storage_dir=storage_dir, + storage_size=storage_size, + num_shard=num_shard, + num_worker=8, + index_endpoint=f"{redis_host}:{redis_port}", + index_prefix=str(redis_prefix), + ) + + data = list(range(num_pages)) + hash_128s = generate_cumulative_hashes(data) + indexer = torch.arange(num_pages, dtype=torch.int32) + + # Write + t = service.create(hash_128s=hash_128s, kv_page_indexer=indexer, mode="w") + _wait_task(t) + _assert_task_finished(t, "write") + + # Query should hit local index + q = service.query(hash_128s) + if not all(q): + raise AssertionError(f"query after write expected all True, got {q}") + + # Readback in same process + kvcache.zero_() + t = service.create(hash_128s=hash_128s, kv_page_indexer=indexer, mode="r") + _wait_task(t) + _assert_task_finished(t, "read") + + if not torch.equal(kvcache, expected): + raise AssertionError("readback mismatch in phase_normal") + + print("[phase_normal] ok") + + +def _phase_restart_verify(storage_dir: Path, redis_host: str, redis_port: int, redis_prefix: str) -> None: + + page_bytes = 4096 + num_pages = 128 + storage_size = 32 * 1024 * 1024 + num_shard = 4 + + # Fresh tensor, deterministic expected + kvcache = torch.zeros((num_pages, page_bytes), dtype=torch.uint8) + expected = _make_kvcache(seed=123, num_pages=num_pages, page_bytes=page_bytes) + + service = _make_service( + kvcache=kvcache, + storage_dir=storage_dir, + storage_size=storage_size, + num_shard=num_shard, + num_worker=8, + index_endpoint=f"{redis_host}:{redis_port}", + index_prefix=str(redis_prefix), + ) + + # If Redis lost its dataset (keys deleted), we need to explicitly recover shard indices + # back to Redis. In multi-node mode this is normally triggered by the coordinator on + # shard ownership changes; here we do it manually. + for shard_id in range(num_shard): + if hasattr(service, "recover_shard_to_redis_smart"): + service.recover_shard_to_redis_smart(int(shard_id)) + elif hasattr(service, "recover_shard_to_redis"): + service.recover_shard_to_redis(int(shard_id)) + + data = list(range(num_pages)) + hash_128s = generate_cumulative_hashes(data) + indexer = torch.arange(num_pages, dtype=torch.int32) + + # On restart, query should already return True (rebuilt from Redis HGETALL) + q = service.query(hash_128s) + if not all(q): + raise AssertionError(f"query after restart expected all True, got {q}") + + # Read from disk + t = service.create(hash_128s=hash_128s, kv_page_indexer=indexer, mode="r") + _wait_task(t) + _assert_task_finished(t, "read after restart") + + if not torch.equal(kvcache, expected): + raise AssertionError("readback mismatch after restart") + + # Append additional writes with new hashes to ensure continued operation. + # Use a disjoint token stream so hashes differ. + kvcache2 = _make_kvcache(seed=456, num_pages=num_pages, page_bytes=page_bytes) + kvcache.copy_(kvcache2) + + data2 = list(range(10_000, 10_000 + num_pages)) + hash_128s2 = generate_cumulative_hashes(data2) + t = service.create(hash_128s=hash_128s2, kv_page_indexer=indexer, mode="w") + _wait_task(t) + _assert_task_finished(t, "append write") + + # Read back appended data + kvcache.zero_() + t = service.create(hash_128s=hash_128s2, kv_page_indexer=indexer, mode="r") + _wait_task(t) + _assert_task_finished(t, "append read") + + if not torch.equal(kvcache, kvcache2): + raise AssertionError("append readback mismatch") + + print("[phase_restart_verify] ok") + + +def _phase_eviction_setup(storage_dir: Path, redis_host: str, redis_port: int, redis_prefix: str) -> None: + """Fill cache beyond capacity to trigger LRU eviction, then verify within-process behavior.""" + + page_bytes = 4096 + pages_per_block = 64 + num_pages = pages_per_block + # Capacity 2 blocks (block_size is 256KB with our params). + storage_size = 512 * 1024 + num_shard = 1 + + indexer = torch.arange(num_pages, dtype=torch.int32) + kvcache = torch.zeros((num_pages, page_bytes), dtype=torch.uint8) + + service = _make_service( + kvcache=kvcache, + storage_dir=storage_dir, + storage_size=storage_size, + num_shard=num_shard, + num_worker=4, + index_endpoint=f"{redis_host}:{redis_port}", + index_prefix=str(redis_prefix), + ) + + def _block_hashes(tag: int) -> list[int]: + data = list(range(tag * 10_000, tag * 10_000 + num_pages)) + return generate_cumulative_hashes(data) + + def _write_block(tag: int, seed: int) -> torch.Tensor: + payload = _make_kvcache(seed=seed, num_pages=num_pages, page_bytes=page_bytes) + kvcache.copy_(payload) + t = service.create(hash_128s=_block_hashes(tag), kv_page_indexer=indexer, mode="w") + _wait_task(t) + _assert_task_finished(t, f"eviction write tag={tag}") + return payload + + # Capacity is 2 blocks. To make the eviction deterministic: + # 1) Write A, then B + # 2) Touch B so it becomes MRU + # 3) Insert C => A should be evicted + expected_a = _write_block(tag=1, seed=111) + expected_b = _write_block(tag=2, seed=222) + + # Touch B to make it MRU. + t = service.create(hash_128s=_block_hashes(2), kv_page_indexer=indexer, mode="r") + _wait_task(t) + _assert_task_finished(t, "eviction read touch B") + if not torch.equal(kvcache, expected_b): + raise AssertionError("eviction: readback of B before eviction mismatch") + + # Insert C, which should evict A. + expected_c = _write_block(tag=3, seed=333) + + # (Optional sanity) Read C back immediately. + kvcache.zero_() + t = service.create(hash_128s=_block_hashes(3), kv_page_indexer=indexer, mode="r") + _wait_task(t) + _assert_task_finished(t, "eviction read C") + if not torch.equal(kvcache, expected_c): + raise AssertionError("eviction: readback of C after insertion mismatch") + + # Now A should be miss, B and C should be hit. + q_a = service.query(_block_hashes(1)) + q_b = service.query(_block_hashes(2)) + q_c = service.query(_block_hashes(3)) + + if any(q_a): + raise AssertionError(f"eviction: expected A to be evicted, got query={q_a}") + if not all(q_b): + raise AssertionError(f"eviction: expected B to exist, got query={q_b}") + if not all(q_c): + raise AssertionError(f"eviction: expected C to exist, got query={q_c}") + + print("[phase_eviction_setup] ok") + + +def _phase_eviction_restart_verify(storage_dir: Path, redis_host: str, redis_port: int, redis_prefix: str) -> None: + """Restart and verify eviction state persists (A missing, B/C readable), and service remains writable.""" + + page_bytes = 4096 + pages_per_block = 64 + num_pages = pages_per_block + storage_size = 512 * 1024 + num_shard = 1 + + indexer = torch.arange(num_pages, dtype=torch.int32) + kvcache = torch.zeros((num_pages, page_bytes), dtype=torch.uint8) + + service = _make_service( + kvcache=kvcache, + storage_dir=storage_dir, + storage_size=storage_size, + num_shard=num_shard, + num_worker=4, + index_endpoint=f"{redis_host}:{redis_port}", + index_prefix=str(redis_prefix), + ) + + def _block_hashes(tag: int) -> list[int]: + data = list(range(tag * 10_000, tag * 10_000 + num_pages)) + return generate_cumulative_hashes(data) + + expected_b = _make_kvcache(seed=222, num_pages=num_pages, page_bytes=page_bytes) + expected_c = _make_kvcache(seed=333, num_pages=num_pages, page_bytes=page_bytes) + + # A should be missing + q_a = service.query(_block_hashes(1)) + if any(q_a): + raise AssertionError(f"eviction restart: expected A miss, got query={q_a}") + + # B/C should be readable + t = service.create(hash_128s=_block_hashes(2), kv_page_indexer=indexer, mode="r") + _wait_task(t) + _assert_task_finished(t, "eviction restart read B") + if not torch.equal(kvcache, expected_b): + raise AssertionError("eviction restart: B mismatch") + + kvcache.zero_() + t = service.create(hash_128s=_block_hashes(3), kv_page_indexer=indexer, mode="r") + _wait_task(t) + _assert_task_finished(t, "eviction restart read C") + if not torch.equal(kvcache, expected_c): + raise AssertionError("eviction restart: C mismatch") + + # Ensure continued operation: write/read D + expected_d = _make_kvcache(seed=444, num_pages=num_pages, page_bytes=page_bytes) + kvcache.copy_(expected_d) + t = service.create(hash_128s=_block_hashes(4), kv_page_indexer=indexer, mode="w") + _wait_task(t) + _assert_task_finished(t, "eviction restart write D") + + kvcache.zero_() + t = service.create(hash_128s=_block_hashes(4), kv_page_indexer=indexer, mode="r") + _wait_task(t) + _assert_task_finished(t, "eviction restart read D") + if not torch.equal(kvcache, expected_d): + raise AssertionError("eviction restart: D mismatch") + + print("[phase_eviction_restart_verify] ok") + + +def _phase_crash_writer(storage_dir: Path, redis_host: str, redis_port: int, redis_prefix: str) -> None: + page_bytes = 4096 + num_pages = 128 + storage_size = 32 * 1024 * 1024 + num_shard = 4 + + # Baseline committed data + kvcache = _make_kvcache(seed=777, num_pages=num_pages, page_bytes=page_bytes) + expected = kvcache.clone() + + service = _make_service( + kvcache=kvcache, + storage_dir=storage_dir, + storage_size=storage_size, + num_shard=num_shard, + num_worker=8, + index_endpoint=f"{redis_host}:{redis_port}", + index_prefix=str(redis_prefix), + ) + + indexer = torch.arange(num_pages, dtype=torch.int32) + + data = list(range(20_000, 20_000 + num_pages)) + hash_128s = generate_cumulative_hashes(data) + + t = service.create(hash_128s=hash_128s, kv_page_indexer=indexer, mode="w") + _wait_task(t) + _assert_task_finished(t, "baseline write") + + # Start another write but do NOT wait; crash immediately. + kvcache.copy_(_make_kvcache(seed=888, num_pages=num_pages, page_bytes=page_bytes)) + data2 = list(range(30_000, 30_000 + num_pages)) + hash_128s2 = generate_cumulative_hashes(data2) + _ = service.create(hash_128s=hash_128s2, kv_page_indexer=indexer, mode="w") + + # Abrupt termination: bypass destructors. + # This simulates SIGKILL / power loss more closely than sys.exit. + print("[phase_crash_writer] crashing now") + os._exit(137) + + +def _phase_crash_verify(storage_dir: Path, redis_host: str, redis_port: int, redis_prefix: str) -> None: + page_bytes = 4096 + num_pages = 128 + storage_size = 32 * 1024 * 1024 + num_shard = 4 + + kvcache = torch.zeros((num_pages, page_bytes), dtype=torch.uint8) + expected = _make_kvcache(seed=777, num_pages=num_pages, page_bytes=page_bytes) + + service = _make_service( + kvcache=kvcache, + storage_dir=storage_dir, + storage_size=storage_size, + num_shard=num_shard, + num_worker=8, + index_endpoint=f"{redis_host}:{redis_port}", + index_prefix=str(redis_prefix), + ) + + indexer = torch.arange(num_pages, dtype=torch.int32) + data = list(range(20_000, 20_000 + num_pages)) + hash_128s = generate_cumulative_hashes(data) + + # Baseline should still be readable after crash. + q = service.query(hash_128s) + if not all(q): + raise AssertionError(f"query after crash expected all True for committed baseline, got {q}") + + t = service.create(hash_128s=hash_128s, kv_page_indexer=indexer, mode="r") + _wait_task(t) + _assert_task_finished(t, "read after crash") + + if not torch.equal(kvcache, expected): + raise AssertionError("baseline readback mismatch after crash") + + # Ensure service remains usable: do another write/read. + kvcache2 = _make_kvcache(seed=999, num_pages=num_pages, page_bytes=page_bytes) + kvcache.copy_(kvcache2) + data3 = list(range(40_000, 40_000 + num_pages)) + hash_128s3 = generate_cumulative_hashes(data3) + + t = service.create(hash_128s=hash_128s3, kv_page_indexer=indexer, mode="w") + _wait_task(t) + _assert_task_finished(t, "post-crash write") + + kvcache.zero_() + t = service.create(hash_128s=hash_128s3, kv_page_indexer=indexer, mode="r") + _wait_task(t) + _assert_task_finished(t, "post-crash read") + + if not torch.equal(kvcache, kvcache2): + raise AssertionError("post-crash readback mismatch") + + print("[phase_crash_verify] ok") + + +def _run_phase_self(args: list[str], *, env: dict[str, str], timeout: int = 120) -> subprocess.CompletedProcess: + return subprocess.run( + [sys.executable, str(Path(__file__).resolve()), *args], + cwd=str(TESTS_DIR), + env=env, + capture_output=False, + text=True, + timeout=timeout, + ) + + +def _orchestrate(ns: argparse.Namespace) -> None: + cleanup_storage = not bool(str(getattr(ns, "storage_dir", "") or "").strip()) + + if bool(getattr(ns, "reuse_services", False)): + host = str(getattr(ns, "redis_host", "") or "").strip() or "127.0.0.1" + port = int(getattr(ns, "redis_port", 0) or 0) + if port <= 0: + raise SystemExit("--reuse-services requires --redis-port") + redis = _RedisHandle("reuse", host, port, lambda: None) + else: + redis = _start_redis_or_skip() + + storage_root_cli = str(getattr(ns, "storage_dir", "") or "").strip() + if storage_root_cli: + storage_root = Path(storage_root_cli).resolve() + storage_root.mkdir(parents=True, exist_ok=True) + else: + storage_root = TESTS_DIR / "cache" + storage_root.mkdir(parents=True, exist_ok=True) + + storage_dir = storage_root / "recovery" + redis_prefix = "lightmem_test_" + uuid.uuid4().hex + + storage_dir_evict = storage_root / "eviction" + redis_prefix_evict = "lightmem_test_evict_" + uuid.uuid4().hex + + base_env = os.environ.copy() + + try: + # Normal shutdown / restart + r1 = _run_phase_self( + [ + "--phase", + "normal", + "--storage", + str(storage_dir), + "--index-endpoint", + f"{redis.host}:{redis.port}", + "--index-prefix", + str(redis_prefix), + ], + env=base_env, + timeout=120, + ) + if r1.returncode != 0: + raise RuntimeError(f"phase normal failed: rc={r1.returncode}") + + # Simulate partial Redis key loss while disk/WAL remains. + # (Full dataset loss recovery isn't guaranteed in all builds/configs.) + _redis_delete_prefix_keys(redis.host, redis.port, prefix=redis_prefix, num_shard=4) + + r2 = _run_phase_self( + [ + "--phase", + "restart_verify", + "--storage", + str(storage_dir), + "--index-endpoint", + f"{redis.host}:{redis.port}", + "--index-prefix", + str(redis_prefix), + ], + env=base_env, + timeout=120, + ) + if r2.returncode != 0: + raise RuntimeError(f"phase restart_verify failed: rc={r2.returncode}") + + # Abnormal exit / restart + r3 = _run_phase_self( + [ + "--phase", + "crash_writer", + "--storage", + str(storage_dir), + "--index-endpoint", + f"{redis.host}:{redis.port}", + "--index-prefix", + str(redis_prefix), + ], + env=base_env, + timeout=120, + ) + if r3.returncode == 0: + raise RuntimeError("phase crash_writer unexpectedly exited cleanly") + + r4 = _run_phase_self( + [ + "--phase", + "crash_verify", + "--storage", + str(storage_dir), + "--index-endpoint", + f"{redis.host}:{redis.port}", + "--index-prefix", + str(redis_prefix), + ], + env=base_env, + timeout=120, + ) + if r4.returncode != 0: + raise RuntimeError(f"phase crash_verify failed: rc={r4.returncode}") + + # LRU eviction correctness across restart + env_evict = base_env.copy() + r5 = _run_phase_self( + [ + "--phase", + "eviction_setup", + "--storage", + str(storage_dir_evict), + "--index-endpoint", + f"{redis.host}:{redis.port}", + "--index-prefix", + str(redis_prefix_evict), + ], + env=env_evict, + timeout=120, + ) + if r5.returncode != 0: + raise RuntimeError(f"phase eviction_setup failed: rc={r5.returncode}") + + r6 = _run_phase_self( + [ + "--phase", + "eviction_restart_verify", + "--storage", + str(storage_dir_evict), + "--index-endpoint", + f"{redis.host}:{redis.port}", + "--index-prefix", + str(redis_prefix_evict), + ], + env=env_evict, + timeout=120, + ) + if r6.returncode != 0: + raise RuntimeError(f"phase eviction_restart_verify failed: rc={r6.returncode}") + + print("✓ Redis + disk recovery robustness test passed") + + finally: + # Best-effort cleanup + if cleanup_storage: + try: + shutil.rmtree(storage_dir, ignore_errors=True) + shutil.rmtree(storage_dir_evict, ignore_errors=True) + except Exception: + pass + redis.cleanup() + + +def main() -> None: + args = _parse_args(sys.argv[1:]) + + if args.phase == "orchestrate": + _orchestrate(args) + return + + storage_dir = Path(args.storage) + if not storage_dir: + raise SystemExit("--storage is required for phase mode") + + endpoint = (args.index_endpoint or "").strip() or "127.0.0.1:6379" + if ":" in endpoint: + host, port_s = endpoint.rsplit(":", 1) + port = int(port_s) + else: + host, port = endpoint, 6379 + prefix = str(args.index_prefix or "lightmem") + + if args.phase == "normal": + _phase_normal(storage_dir, host, port, prefix) + elif args.phase == "restart_verify": + _phase_restart_verify(storage_dir, host, port, prefix) + elif args.phase == "crash_writer": + _phase_crash_writer(storage_dir, host, port, prefix) + elif args.phase == "crash_verify": + _phase_crash_verify(storage_dir, host, port, prefix) + elif args.phase == "eviction_setup": + _phase_eviction_setup(storage_dir, host, port, prefix) + elif args.phase == "eviction_restart_verify": + _phase_eviction_restart_verify(storage_dir, host, port, prefix) + else: + raise SystemExit(f"unknown phase: {args.phase}") + + +if __name__ == "__main__": + main() diff --git a/tests/multi_node/test_13_concurrent_init_shared_dir.py b/tests/multi_node/test_13_concurrent_init_shared_dir.py new file mode 100755 index 0000000..d5168e3 --- /dev/null +++ b/tests/multi_node/test_13_concurrent_init_shared_dir.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 +"""Concurrent init against a shared storage directory. + +Regression test for TOCTOU-style races in LocalStorageEngine::createOrOpenFiles when multiple +services start at the same time pointing to the same `--storage-dir`. + +Expected behavior after the fix: +- Exactly one process becomes the initializer (creates dirs/files and writes `.initialized`). +- Other processes wait for `.initialized` and only open existing files (no create/truncate). + +Standalone or run_all with --reuse-services. +""" + +from __future__ import annotations + +import resource +import subprocess +import sys +import time +from pathlib import Path + +# Allow running this test from repo root where the script dir isn't on sys.path. +_THIS_DIR = Path(__file__).resolve().parent +if str(_THIS_DIR) not in sys.path: + sys.path.insert(0, str(_THIS_DIR)) + +from harness import parse_common_args, require_ports, start_cluster_env + + +WORKER = Path(__file__).with_name("worker_node_ops.py") + + +def _wait_path_exists(path: Path, timeout_s: float) -> None: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + if path.exists(): + return + time.sleep(0.01) + raise TimeoutError(f"timeout waiting for {path}") + + +def _tail(p: subprocess.Popen[str] | None, *, max_chars: int = 4000) -> str: + try: + if p is None or p.stdout is None: + return "" + return (p.stdout.read() or "")[-int(max_chars) :] + except Exception: + return "" + + +def _choose_num_shard() -> int: + # Each shard opens 2 fds (data/meta) per process. Add a conservative reserve for + # stdio, sockets, dylibs, etc. + try: + soft, _hard = resource.getrlimit(resource.RLIMIT_NOFILE) + if soft == resource.RLIM_INFINITY: + soft = 4096 + soft = int(soft) + except Exception: + soft = 1024 + + reserve = 256 + per_shard_fds = 2 + max_by_limit = max(32, (max(0, soft - reserve) // per_shard_fds)) + + # We don't need huge scale here; we only need concurrent init ordering. + return int(max(64, min(256, max_by_limit))) + + +def _wait_marker_or_fail( + *, + marker: Path, + procs: dict[str, subprocess.Popen[str] | None], + timeout_s: float, +) -> None: + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + if marker.exists(): + return + for name, p in procs.items(): + if p is None: + continue + rc = p.poll() + if rc is not None and rc != 0: + out = _tail(p) + raise RuntimeError(f"worker {name} exited rc={rc}\n{out}") + time.sleep(0.05) + raise TimeoutError(f"timeout waiting for {marker}") + + +def main(argv: list[str] | None = None) -> int: + ns = parse_common_args(argv) + redis_host, redis_port, etcd_host, etcd_port = require_ports(ns) + + num_shard = _choose_num_shard() + storage_size = 24 * 1024 * 1024 * 1024 # sparse file preallocation + + try: + soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) + except Exception: + soft, hard = -1, -1 + print(f"[test] RLIMIT_NOFILE soft={soft} hard={hard} num_shard={num_shard}", flush=True) + + env = start_cluster_env( + reuse_services=bool(ns.reuse_services), + redis_host=redis_host, + redis_port=redis_port, + etcd_host=etcd_host, + etcd_port=etcd_port, + storage_dir=(ns.storage_dir or None), + cleanup_storage=not bool(ns.storage_dir), + ) + + p0: subprocess.Popen[str] | None = None + p1: subprocess.Popen[str] | None = None + try: + marker_dir = Path(env.storage_dir) / "_markers" + marker_dir.mkdir(parents=True, exist_ok=True) + + # These are created by LocalStorageEngine::createOrOpenFiles. + init_lock_dir = Path(str(env.storage_dir) + ".init.lock") + init_marker = Path(str(env.storage_dir) + ".initialized") + + node0_started = marker_dir / "node0_started" + node1_started = marker_dir / "node1_started" + + # Start node-0 first; it should win the global init lock in most cases. + p0 = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "1", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-0", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--op", + "idle", + "--duration-sec", + "6", + "--started-file", + str(node0_started), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + # Wait until we observe the lock dir (or marker, if init is extremely fast). + deadline = time.time() + 10.0 + while time.time() < deadline: + if init_lock_dir.exists() or init_marker.exists(): + break + time.sleep(0.01) + + p1 = subprocess.Popen( + [ + sys.executable, + str(WORKER), + "--storage-dir", + str(env.storage_dir), + "--storage-size", + str(storage_size), + "--num-shard", + str(num_shard), + "--num-worker", + "1", + "--redis", + f"{env.redis.host}:{env.redis.port}", + "--etcd", + f"{env.etcd.host}:{env.etcd.port}", + "--node-id", + "node-1", + "--ttl", + "6", + "--reconcile-sec", + "1.0", + "--op", + "idle", + "--duration-sec", + "6", + "--started-file", + str(node1_started), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + procs = {"node-0": p0, "node-1": p1} + + # Regression check: + # If the follower reaches "started" before the marker exists, it means it didn't wait. + # (If init is very fast, the marker may already exist and this is fine.) + deadline = time.time() + 5.0 + while time.time() < deadline: + if node1_started.exists() and not init_marker.exists(): + raise AssertionError("follower started before init marker existed") + if init_marker.exists(): + break + # Also fail-fast if a worker exits. + for name, p in procs.items(): + if p is None: + continue + rc = p.poll() + if rc is not None and rc != 0: + raise RuntimeError(f"worker {name} exited rc={rc}\n{_tail(p)}") + time.sleep(0.02) + + _wait_marker_or_fail(marker=init_marker, procs=procs, timeout_s=90.0) + _wait_path_exists(node0_started, timeout_s=30.0) + _wait_path_exists(node1_started, timeout_s=30.0) + + # After completion, lock dir should be released (best-effort rmdir). + assert not init_lock_dir.exists(), f"init lock dir still exists: {init_lock_dir}" + + # Spot-check a few shards for expected files. + sample_sids = [0] + if num_shard > 1: + sample_sids.append(1) + if num_shard > 2: + sample_sids.append(num_shard - 1) + for sid in sample_sids: + shard_dir = Path(f"{env.storage_dir}_{sid}") + assert shard_dir.exists() and shard_dir.is_dir(), f"missing shard dir {shard_dir}" + for name in ("data", "meta"): + p = shard_dir / name + assert p.exists(), f"missing {p}" + + return 0 + finally: + for p in (p0, p1): + if p is None: + continue + p.terminate() + try: + p.wait(timeout=5) + except subprocess.TimeoutExpired: + p.kill() + p.wait(timeout=5) + + # Print tails to aid debugging on failures. + for name, p in (("node-0", p0), ("node-1", p1)): + out = _tail(p) + if out: + print(f"[{name} output tail]\n" + out) + + env.cleanup() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/worker_crash_node.py b/tests/multi_node/worker_crash_node.py new file mode 100755 index 0000000..f8d3cc3 --- /dev/null +++ b/tests/multi_node/worker_crash_node.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 +"""Worker process for crash recovery test. + +It uses env vars set by the parent test: +- LIGHTMEM_TEST_REDIS=host:port +- LIGHTMEM_TEST_ETCD=host:port +- LIGHTMEM_TEST_STORAGE=/path/to/shared/dir + +Modes: +- write_then_crash: start service, write a few hashes, then os._exit(137) +- restart_and_verify: start service, query+read those hashes, exit 0 +""" + +from __future__ import annotations + +import argparse +import os +import time + +import torch + +from light_mem import PyLocalCacheService + +from harness import build_hash_128s_for_blocks, iter_hash_ids + + +def _make_kvcache(*, seed: int, num_pages: int, page_bytes: int) -> torch.Tensor: + torch.manual_seed(seed) + return torch.randint(0, 256, size=(num_pages, page_bytes), dtype=torch.uint8, device="cpu") + + +def _wait_task(task, timeout_s: float = 60.0) -> None: + deadline = time.time() + timeout_s + while not task.ready(): + if time.time() > deadline: + raise TimeoutError("task timeout") + time.sleep(0.001) + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--mode", required=True, choices=["write_then_crash", "restart_and_verify"]) + args = p.parse_args(argv) + + redis_ep = os.environ["LIGHTMEM_TEST_REDIS"] + etcd_ep = os.environ["LIGHTMEM_TEST_ETCD"] + storage = os.environ["LIGHTMEM_TEST_STORAGE"] + + num_shard = 192 + storage_size = 6 * 1024 * 1024 * 1024 + + page_bytes = 4096 + num_pages = 2048 + + kvcache = _make_kvcache(seed=4242, num_pages=num_pages, page_bytes=page_bytes) + + svc = PyLocalCacheService( + kvcache_tensor=kvcache, + file=str(storage), + storage_size=storage_size, + num_shard=num_shard, + num_worker=2, + index_endpoint=redis_ep, + coord_endpoints=etcd_ep, + coord_node_id="crash-node", + coord_ttl=6, + ) + + block_size = int(svc._c.block_size()) + n_pages = block_size // page_bytes + indexer = torch.arange(n_pages, dtype=torch.int32) + + hashes = iter_hash_ids(count=8, seed=5150) + + if args.mode == "write_then_crash": + for hid in hashes: + h128s = build_hash_128s_for_blocks(block_hash_ids=[hid], pages_per_block=n_pages) + t = svc.create(hash_128s=h128s, kv_page_indexer=indexer, mode="w") + _wait_task(t) + + # Abrupt exit (simulate crash) + os._exit(137) + + # restart_and_verify + for hid in hashes: + assert svc.query([hid]) == [True] + h128s = build_hash_128s_for_blocks(block_hash_ids=[hid], pages_per_block=n_pages) + t = svc.create(hash_128s=h128s, kv_page_indexer=indexer, mode="r") + _wait_task(t) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/multi_node/worker_node_ops.py b/tests/multi_node/worker_node_ops.py new file mode 100755 index 0000000..bd1d143 --- /dev/null +++ b/tests/multi_node/worker_node_ops.py @@ -0,0 +1,1494 @@ +#!/usr/bin/env python3 +"""Helper process for multi-node tests. + +Why this exists: +- Some environments crash (segfault) when multiple PyLocalCacheService instances run + in the same Python process and execute create/read tasks. +- Spawning one service per process better matches real multi-node deployments and + avoids shared in-process state. + +This worker starts exactly one PyLocalCacheService and optionally performs a single +operation (write/read/query/recover/stress). + +Exit code 0 means success. +""" + +from __future__ import annotations + +import argparse +import json +import os +import random +import socket +import signal +import sys +import time +import hashlib +from pathlib import Path +from collections import deque + +import torch + +from light_mem import PyLocalCacheService, PyState + + +def _scan_etcd_shard_owners(*, etcd_endpoint: str, prefix: str, num_shards: int) -> dict[int, str]: + try: + from harness import etcd_client # type: ignore + + host, port_s = str(etcd_endpoint).rsplit(":", 1) + cli = etcd_client(str(host), int(port_s)) + base = f"{prefix.rstrip('/')}/shards/" + owners: dict[int, str] = {} + for value, meta in cli.get_prefix(base): + try: + key = meta.key.decode("utf-8") + except Exception: + continue + if not key.endswith("/owner"): + continue + try: + sid_str = key[len(base) :].split("/", 1)[0] + sid = int(sid_str) + except Exception: + continue + if 0 <= sid < int(num_shards) and value is not None: + try: + owners[sid] = value.decode("utf-8") + except Exception: + continue + return owners + except Exception: + return {} + + +def _maybe_wait_for_single_node_ownership( + *, + etcd_endpoint: str, + index_prefix: str, + node_id: str, + num_shards: int, + storage_size_bytes: int, +) -> None: + """Avoid starting IO before etcd ownership converges. + + In distributed/shared-storage mode we intentionally fail-closed (no writable shards) + right after service creation until the coordinator thread applies assignments. + If we start read/write immediately, tasks can be aborted and benchmarks become misleading. + + - If LIGHTMEM_EXPECT_OWN_ALL_SHARDS=1: require this node owns all shards. + - Otherwise (multi-node): wait until all shards have an owner AND this node owns >=1 shard. + """ + ep = (etcd_endpoint or "").strip() + if not ep: + return + + prefix = (index_prefix or "").strip() + expect_all = str(os.environ.get("LIGHTMEM_EXPECT_OWN_ALL_SHARDS") or "").strip() == "1" + + deadline = time.time() + (90.0 if expect_all else 30.0) + last_msg = 0.0 + while True: + owners = _scan_etcd_shard_owners(etcd_endpoint=ep, prefix=prefix, num_shards=int(num_shards)) + owned = sum(1 for o in owners.values() if o == str(node_id)) + total_owned_keys = int(len(owners)) + uniq_nodes = len(set(owners.values())) if owners else 0 + + # Print at most once per second. + now = time.time() + if now - last_msg >= 1.0: + total_gb = float(int(storage_size_bytes)) / (1024.0**3) + eff_gb = total_gb * (float(owned) / float(max(1, int(num_shards)))) + print( + f"[coord] etcd_prefix={prefix} owned_shards={owned}/{int(num_shards)} " + f"owners_keys={total_owned_keys}/{int(num_shards)} uniq_nodes={uniq_nodes} " + f"effective_capacity≈{eff_gb:.2f}GB (of {total_gb:.2f}GB)", + flush=True, + ) + last_msg = now + + if expect_all: + ready = (owned >= int(num_shards)) and (total_owned_keys >= int(num_shards)) + else: + ready = (total_owned_keys >= int(num_shards)) and (owned >= 1) + + if ready: + return + + if time.time() >= deadline: + if expect_all: + raise RuntimeError( + "single-node run expected to own all shards, but ownership did not converge. " + f"prefix={prefix} owned={owned}/{int(num_shards)} owners_keys={total_owned_keys}/{int(num_shards)}; " + "this usually means other nodes are still registered under the same etcd prefix. " + "Use a unique --coord-prefix (recommended) or stop other nodes." + ) + raise RuntimeError( + "multi-node run expected shard ownership to converge, but it did not. " + f"prefix={prefix} owned={owned}/{int(num_shards)} owners_keys={total_owned_keys}/{int(num_shards)}." + ) + + time.sleep(0.2) + + +def _estimate_effective_capacity_blocks( + *, + etcd_endpoint: str, + node_id: str, + prefix: str, + num_shards: int, + capacity_blocks_total: int, +) -> tuple[int, int, int]: + """Estimate this node's effective capacity in blocks. + + In etcd coordinated mode, writes are only allowed to this node's owned shards. + Because eviction is per-shard, a node's usable capacity is proportional to + owned_shards / num_shards. + + Returns (effective_capacity_blocks, owned_shards, uniq_nodes). + """ + try: + owners = _scan_etcd_shard_owners(etcd_endpoint=str(etcd_endpoint), prefix=str(prefix), num_shards=int(num_shards)) + owned = sum(1 for o in owners.values() if o == str(node_id)) + uniq_nodes = len(set(owners.values())) if owners else 0 + eff = int(int(capacity_blocks_total) * (float(owned) / float(max(1, int(num_shards))))) + return max(1, int(eff)), int(owned), int(uniq_nodes) + except Exception: + return max(1, int(capacity_blocks_total)), 0, 0 + + +def _sha256_bytes(b: bytes) -> str: + try: + return hashlib.sha256(b).hexdigest() + except Exception: + return "" + + +def _tensor_fingerprint_u8(t: torch.Tensor, *, max_bytes: int = 64) -> dict: + # Returns small, stable identifiers for debugging. + try: + cpu = t.detach().to(device="cpu") + if cpu.dtype != torch.uint8: + cpu = cpu.to(dtype=torch.uint8) + raw = cpu.contiguous().numpy().tobytes() + return { + "shape": list(cpu.shape), + "dtype": str(cpu.dtype), + "nbytes": int(len(raw)), + "sha256": _sha256_bytes(raw), + "head_hex": raw[: int(max_bytes)].hex(), + "tail_hex": raw[-int(max_bytes) :].hex() if len(raw) > int(max_bytes) else "", + } + except Exception as e: + return {"error": repr(e)} + + +def _best_effort_etcd_owner(*, etcd_endpoint: str, prefix: str, shard_id: int) -> str | None: + # Optional: avoid hard dependency; only used for debug dumps. + try: + # worker runs from tests/multi_node, so this import should work. + from harness import etcd_client # type: ignore + + host, port_s = str(etcd_endpoint).rsplit(":", 1) + cli = etcd_client(str(host), int(port_s)) + key = f"{prefix.rstrip('/')}/shards/{int(shard_id)}/owner" + val = cli.get(key) + # etcd_client.get may return either bytes or (bytes|None, meta|None). + if isinstance(val, tuple) and len(val) >= 1: + val = val[0] + if not val: + return None + try: + return val.decode("utf-8") + except Exception: + return str(val) + except Exception: + return None + + +def _debug_dump(*, path: str, obj: dict) -> None: + # Always print a compact line to stdout (orchestrator tails stdout). + try: + print("[debug_dump] " + json.dumps(obj, sort_keys=True), flush=True) + except Exception: + pass + + if not path: + return + p = Path(path) + try: + p.parent.mkdir(parents=True, exist_ok=True) + # Write one json per dump (atomic replace) for easy inspection. + tmp = p.with_suffix(p.suffix + ".tmp") + tmp.write_text(json.dumps(obj, sort_keys=True, indent=2), encoding="utf-8") + tmp.replace(p) + except Exception: + pass + + +def _touch(path: str) -> None: + if not path: + return + p = Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text("ok", encoding="utf-8") + + +def _wait_file(path: str, timeout_s: float) -> None: + if not path: + return + p = Path(path) + deadline = time.time() + float(timeout_s) + while time.time() < deadline: + if p.exists(): + return + time.sleep(0.01) + raise TimeoutError(f"timeout waiting for {path}") + + +def _make_kvcache(*, seed: int, num_pages: int, page_bytes: int) -> torch.Tensor: + torch.manual_seed(int(seed)) + return torch.randint(0, 256, size=(int(num_pages), int(page_bytes)), dtype=torch.uint8, device="cpu") + + +def _hash_128s_for_one_block(*, hash_id: int, pages_per_block: int) -> list[int]: + # Any values work for the first (pages_per_block-1) entries; PyLocalCacheService + # uses the last element of each block as the block hash. + dummy = 1 + out: list[int] = [] + for _ in range(int(pages_per_block) - 1): + out.append(dummy) + dummy += 1 + out.append(int(hash_id)) + return out + + +def _wait_task(task, timeout_s: float = 60.0) -> None: + deadline = time.time() + float(timeout_s) + while not task.ready(): + if time.time() > deadline: + raise TimeoutError("task timeout") + time.sleep(0.001) + + # Surface failures early; otherwise benchmarks can report fake bandwidth. + try: + st = task.state() + if any(s == PyState.Aborted for s in st): + raise RuntimeError(f"task aborted: {st}") + except Exception: + # Best-effort: some bindings may not expose state reliably. + pass + + +def _atomic_write_json(path: str, obj: dict) -> None: + if not path: + return + p = Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + tmp = p.with_suffix(p.suffix + ".tmp") + tmp.write_text(json.dumps(obj, sort_keys=True), encoding="utf-8") + tmp.replace(p) + + +def _hostname() -> str: + try: + return socket.gethostname() + except Exception: + return "unknown" + + +def _append_jsonl(path: str, obj: dict) -> None: + if not path: + return + p = Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + line = json.dumps(obj, sort_keys=True) + # Best-effort append. On some remote filesystems, atomic append is not guaranteed; + # readers are expected to tolerate partial/garbled lines. + with p.open("a", encoding="utf-8") as f: + f.write(line + "\n") + f.flush() + + +def _read_jsonl_tail(path: str, *, max_bytes: int = 256 * 1024) -> list[dict]: + if not path: + return [] + p = Path(path) + if not p.exists(): + return [] + try: + with p.open("rb") as f: + try: + f.seek(0, os.SEEK_END) + end = f.tell() + start = max(0, end - int(max_bytes)) + f.seek(start, os.SEEK_SET) + raw = f.read() + except Exception: + raw = p.read_bytes() + except Exception: + return [] + + # If we started from the middle of a line, drop the first partial line. + try: + text = raw.decode("utf-8", errors="replace") + except Exception: + return [] + lines = text.splitlines() + if not lines: + return [] + if len(raw) >= max_bytes and lines: + lines = lines[1:] + + out: list[dict] = [] + for ln in lines: + ln = ln.strip() + if not ln: + continue + try: + obj = json.loads(ln) + except Exception: + continue + if isinstance(obj, dict): + out.append(obj) + return out + + +def _expected_block(*, hash_id: int, pages_per_block: int, page_bytes: int) -> torch.Tensor: + # Deterministic across processes/machines: seed is derived solely from hash_id. + # This allows any node to verify a block written by any other node. + g = torch.Generator(device="cpu") + g.manual_seed(int(hash_id) & 0xFFFFFFFFFFFFFFFF) + return torch.randint( + 0, + 256, + size=(int(pages_per_block), int(page_bytes)), + dtype=torch.uint8, + device="cpu", + generator=g, + ) + + +def _write_one_verify( + *, + svc: PyLocalCacheService, + hash_id: int, + index_prefix: str = "lightmem", + pages_per_block: int, + page_bytes: int, + indexer: torch.Tensor, + kvcache: torch.Tensor, + node_id: str | None = None, + redis_endpoint: str | None = None, + etcd_endpoint: str | None = None, + num_shard: int | None = None, + debug_dump_file: str = "", +) -> None: + # In distributed mode, writes may be temporarily skipped until shard assignments + # are loaded (LocalStorageEngine treats all shards as draining by default). + # Also, dedupe may skip writes for already-existing hashes. + # For a verification write, we require that the hash becomes queryable and + # that an immediate read returns the exact expected bytes. + expected = _expected_block(hash_id=hash_id, pages_per_block=pages_per_block, page_bytes=page_bytes) + h128s = _hash_128s_for_one_block(hash_id=hash_id, pages_per_block=pages_per_block) + + deadline = time.time() + 60.0 + last_read_state = None + last_write_state = None + attempts = 0 + while time.time() < deadline: + attempts += 1 + kvcache[indexer] = expected + t = svc.create(hash_128s=h128s, kv_page_indexer=indexer, mode="w") + _wait_task(t, timeout_s=60.0) + + try: + last_write_state = t.state() + except Exception: + last_write_state = None + + # Ensure it is actually persisted/indexed. + if svc.query([int(hash_id)]) != [True]: + time.sleep(0.05) + continue + + # Clear then read back and compare. + kvcache[indexer] = 0 + t2 = svc.create(hash_128s=h128s, kv_page_indexer=indexer, mode="r") + _wait_task(t2, timeout_s=60.0) + st = None + try: + st = t2.state() + last_read_state = st + except Exception: + pass + + # If read aborted transiently (e.g., mapping not fully converged), retry. + if st is None: + time.sleep(0.01) + continue + if any(s != PyState.Finished for s in st): + time.sleep(0.05) + continue + + got = kvcache[indexer].cpu() + if got.shape == expected.shape and torch.equal(got, expected): + return + + # If hash exists but content mismatches, that's a hard error. + shard_guess = None + try: + if num_shard is not None and int(num_shard) > 0: + shard_guess = int(hash_id) % int(num_shard) + except Exception: + shard_guess = None + + owner = None + if etcd_endpoint and shard_guess is not None: + prefix = (index_prefix or "").strip() + owner = _best_effort_etcd_owner(etcd_endpoint=str(etcd_endpoint), prefix=prefix, shard_id=int(shard_guess)) + + dump = { + "event": "write_verify_mismatch", + "ts": float(time.time()), + "node_id": str(node_id or ""), + "pid": int(os.getpid()), + "host": _hostname(), + "hash_id": int(hash_id), + "attempts": int(attempts), + "pages_per_block": int(pages_per_block), + "page_bytes": int(page_bytes), + "block_size": int(pages_per_block) * int(page_bytes), + "num_shard": int(num_shard) if num_shard is not None else None, + "shard_guess_mod": shard_guess, + "shard_owner": owner, + "redis": str(redis_endpoint or ""), + "etcd": str(etcd_endpoint or ""), + "python": str(sys.version).split("\n", 1)[0], + "torch": getattr(torch, "__version__", ""), + "query_result": None, + "write_state": [str(x) for x in (last_write_state or [])] if last_write_state is not None else None, + "read_state": [str(x) for x in (last_read_state or [])] if last_read_state is not None else None, + "expected": _tensor_fingerprint_u8(expected), + "got": _tensor_fingerprint_u8(got), + } + try: + dump["query_result"] = svc.query([int(hash_id)]) + except Exception as e: + dump["query_result"] = {"error": repr(e)} + + _debug_dump(path=debug_dump_file, obj=dump) + raise AssertionError(f"data mismatch for hash_id={hash_id}") + + raise TimeoutError(f"timeout waiting verified write to become readable for hash_id={hash_id}, last_read_state={last_read_state}") + + +def _read_one_verify(*, svc: PyLocalCacheService, hash_id: int, pages_per_block: int, page_bytes: int, indexer: torch.Tensor, kvcache: torch.Tensor) -> bool: + # Returns True if present+verified, False if missing. + # + # NOTE: In etcd distributed mode, LocalStorageEngine::queryMany gates hits by shard ownership + # (isShardOwned). That means a non-owner node may see query()==False even though the block is + # readable via Redis global index + shared storage. + # For cross-node verification we must attempt a real read and validate bytes. + expected = _expected_block(hash_id=hash_id, pages_per_block=pages_per_block, page_bytes=page_bytes) + h128s = _hash_128s_for_one_block(hash_id=hash_id, pages_per_block=pages_per_block) + # Tolerate very short transient failures (e.g., mapping convergence). + for _ in range(3): + kvcache[indexer] = 0 + t = svc.create(hash_128s=h128s, kv_page_indexer=indexer, mode="r") + _wait_task(t, timeout_s=60.0) + try: + st = t.state() + if any(s != PyState.Finished for s in st): + time.sleep(0.05) + continue + except Exception: + pass + + got = kvcache[indexer].cpu() + if got.shape != expected.shape or not torch.equal(got, expected): + raise AssertionError(f"data mismatch for hash_id={hash_id}") + return True + + return False +def _query_one_block_exists(*, svc: PyLocalCacheService, hash_id: int, pages_per_block: int) -> bool: + """Best-effort existence query for a single block. + + IMPORTANT: PyLocalCacheService.query expects a per-page hash_128s list, grouped by block. + Passing a single int will not represent one full block and can return false negatives. + + In etcd distributed mode, C++ queryMany answers "readable on this node" (ownership-gated). + For eviction probing within a node's owned shards, this is acceptable as a lightweight signal. + """ + try: + h128s = _hash_128s_for_one_block(hash_id=int(hash_id), pages_per_block=int(pages_per_block)) + ok = svc.query(h128s) + if isinstance(ok, list) and ok: + return bool(ok[0]) + return False + except Exception: + return False + + +def _true_lru_eviction(*, svc: PyLocalCacheService) -> tuple[bool, int]: + """Return (eviction_observed, eviction_count) from the core service. + + This is the only reliable signal: it is raised when LocalCacheIndex actually + evicts a victim due to capacity pressure. + """ + # Prefer the Python wrapper methods (PyLocalCacheService) if present. + try: + cnt = int(getattr(svc, "eviction_count")()) + obs = bool(getattr(svc, "eviction_observed")()) + return obs, cnt + except Exception: + pass + + # Fallback: call the bound C++ service directly. + try: + c = getattr(svc, "_c", None) + if c is None: + return False, 0 + cnt = int(getattr(c, "eviction_count")()) + if hasattr(c, "eviction_observed"): + obs = bool(getattr(c, "eviction_observed")()) + else: + obs = cnt > 0 + return obs, cnt + except Exception: + return False, 0 + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--storage-dir", required=True) + p.add_argument("--storage-size", type=int, required=True) + p.add_argument("--num-shard", type=int, required=True) + p.add_argument("--num-worker", type=int, default=1) + + p.add_argument( + "--disable-coord", + action="store_true", + help="disable Redis/Etcd coordination and index backend (standalone local mode; bandwidth-only).", + ) + + # Endpoints are required only when coordination is enabled. + p.add_argument("--redis", default="", help="host:port") + p.add_argument("--etcd", default="", help="host:port") + p.add_argument( + "--index-prefix", + default="lightmem", + help="isolate etcd/redis namespaces (default: lightmem)", + ) + + p.add_argument("--node-id", required=True) + p.add_argument("--ttl", type=int, default=6) + p.add_argument("--reconcile-sec", type=float, default=1.0) + + p.add_argument("--num-pages", type=int, default=2048) + p.add_argument("--page-bytes", type=int, default=4096) + + p.add_argument( + "--op", + required=True, + choices=[ + "idle", + "write", + "read", + "query", + "recover", + "stress_write", + "stress_read", + "bench_write", + "bench_read", + "write_wait_recover", + "write_verify", + "read_verify", + "loop_rw_verify", + ], + ) + p.add_argument( + "--force-redis-read", + action="store_true", + help="force Redis resolution path for read/read_verify by dropping local shard ownership", + ) + p.add_argument("--hash-id", default="", help="integer or 0x-prefixed hex") + p.add_argument( + "--hash-ids", + default="", + help="comma-separated hash ids (ints or 0x...); overrides --hash-id when set", + ) + p.add_argument("--duration-sec", type=float, default=5.0) + p.add_argument("--hold-sec", type=float, default=0.0, help="sleep after finishing op") + + p.add_argument("--stats-file", default="") + p.add_argument("--report-interval-sec", type=float, default=2.0) + p.add_argument("--ready-file", default="", help="touch once first verified rw succeeds") + p.add_argument("--resident-window", type=int, default=64, help="recent keys expected to remain readable") + p.add_argument("--evict-probe-gap", type=int, default=1024, help="probe key this many writes behind to confirm LRU") + + # Cross-node probe verification. + p.add_argument("--probe-file", default="", help="shared jsonl file for cross-node read verification") + p.add_argument("--hot-set-size", type=int, default=8, help="number of hot keys to keep refreshed/published") + p.add_argument("--probe-publish-interval-sec", type=float, default=2.0) + p.add_argument("--probe-read-interval-sec", type=float, default=3.0) + p.add_argument("--probe-ttl-sec", type=float, default=20.0, help="probe entries older than this are ignored") + p.add_argument("--probe-max-read", type=int, default=8, help="max probe keys to verify each interval") + + # High-throughput bench. + p.add_argument( + "--batch-blocks", + type=int, + default=64, + help="for bench_write: number of blocks to write per create() (capped by num_pages/pages_per_block)", + ) + p.add_argument( + "--read-window-blocks", + type=int, + default=2048, + help="for bench_read: read within [base, base+window) repeatedly to avoid missing keys", + ) + + p.add_argument( + "--debug-dump-file", + default="", + help="optional path to write a JSON debug dump when verification fails", + ) + + p.add_argument("--started-file", default="") + p.add_argument("--phase1-file", default="", help="optional marker after phase-1 completes") + p.add_argument("--trigger-file", default="", help="for write_wait_recover: wait until this file exists") + p.add_argument("--done-file", default="") + + args = p.parse_args(argv) + + disable_coord = bool(getattr(args, "disable_coord", False)) + + index_prefix = (str(getattr(args, "index_prefix", "")) or "").strip() + + redis = str(args.redis or "") + etcd = str(args.etcd or "") + if not disable_coord: + if not redis or not etcd: + raise ValueError("--redis and --etcd are required unless --disable-coord is set") + else: + # Ensure we never attempt to talk to external services. + redis = "" + etcd = "" + + kvcache = _make_kvcache(seed=hash(args.node_id) & 0xFFFFFFFF, num_pages=args.num_pages, page_bytes=args.page_bytes) + + svc = PyLocalCacheService( + kvcache_tensor=kvcache, + file=str(args.storage_dir), + storage_size=int(args.storage_size), + num_shard=int(args.num_shard), + num_worker=int(args.num_worker), + index_endpoint=("" if disable_coord else redis), + index_prefix=index_prefix, + coord_endpoints=("" if disable_coord else etcd), + coord_node_id=str(args.node_id), + coord_ttl=int(args.ttl), + coord_reconcile_sec=float(args.reconcile_sec), + bandwidth_log=False, + ) + + # IMPORTANT: In distributed / shared-storage mode, multiple processes can point to the same + # underlying shard files. LocalStorageEngine defaults to "all shards writable" on startup, + # and only later receives shard assignments from the etcd coordinator thread. + # If we start writes immediately, different nodes can concurrently write the same shard files + # and corrupt data (each process maintains its own in-memory LRU index). + # Mitigation for tests: force all shards to non-writable+draining until coordinator updates. + if not disable_coord: + try: + if etcd: + svc._c.update_shard_assignments([], [], []) + except Exception: + pass + + if bool(getattr(args, "force_redis_read", False)) and not disable_coord: + try: + svc._c.update_shard_assignments([], [], []) + except Exception: + pass + + _touch(args.started_file) + + def _sigterm(_signum, _frame): + raise SystemExit(0) + + signal.signal(signal.SIGTERM, _sigterm) + + def _parse_one(s: str) -> int: + s = str(s).strip().lower() + if not s: + raise ValueError("empty hash id") + if s.startswith("0x"): + return int(s, 16) + return int(s) + + hash_ids: list[int] = [] + if args.hash_ids: + parts = [p.strip() for p in str(args.hash_ids).split(",") if p.strip()] + hash_ids = [_parse_one(x) for x in parts] + elif args.hash_id: + hash_ids = [_parse_one(args.hash_id)] + + try: + if args.op == "idle": + time.sleep(float(args.duration_sec)) + + elif args.op == "query": + if not hash_ids: + raise ValueError("--hash-id or --hash-ids required") + # NOTE: PyLocalCacheService.query expects a cumulative hash_128s list + # (grouped by pages_per_block), and _hash() will only take the last + # element of each block. The multi_node tests typically operate on + # per-block ids (one 128-bit int per block). For correctness, query + # each id separately. + for hid in hash_ids: + ok = svc.query([hid]) + if ok != [True]: + raise AssertionError(f"query failed for {hid}: {ok}") + + elif args.op in ("write", "read"): + if not hash_ids: + raise ValueError("--hash-id or --hash-ids required") + + # In distributed/shared-storage mode we fail-closed at startup (no writable shards) + # until the coordinator thread applies shard assignments. + # For plain "write" ops (used by bootstrap tests), we must wait for ownership to + # converge; otherwise writes can be effectively skipped and Redis global index + # mappings will never be published. + if args.op == "write" and not disable_coord: + _maybe_wait_for_single_node_ownership( + etcd_endpoint=etcd, + index_prefix=index_prefix, + node_id=str(args.node_id), + num_shards=int(args.num_shard), + storage_size_bytes=int(args.storage_size), + ) + + block_size = int(svc._c.block_size()) + pages_per_block = block_size // int(args.page_bytes) + indexer = torch.arange(pages_per_block, dtype=torch.int32) + mode = "w" if args.op == "write" else "r" + for hid in hash_ids: + h128s = _hash_128s_for_one_block(hash_id=hid, pages_per_block=pages_per_block) + task = svc.create(hash_128s=h128s, kv_page_indexer=indexer, mode=mode) + _wait_task(task, timeout_s=60.0) + + # Ensure the Redis global mapping is actually visible before reporting success. + # This prevents tests from racing on the index publish path (journal worker). + if args.op == "write" and not disable_coord: + field_hex = format(int(hid), "032x") + deadline = time.time() + 10.0 + while time.time() < deadline: + try: + ok = svc._c.query([field_hex]) + if ok == [True]: + break + except Exception: + pass + time.sleep(0.02) + else: + raise TimeoutError(f"timeout waiting Redis global index mapping for {field_hex}") + + elif args.op in ("write_verify", "read_verify"): + if not hash_ids: + raise ValueError("--hash-id or --hash-ids required") + block_size = int(svc._c.block_size()) + pages_per_block = block_size // int(args.page_bytes) + indexer = torch.arange(pages_per_block, dtype=torch.int32) + + # One-shot verification. Note: read_verify tolerates missing keys (returns False) + # but will raise if present data mismatches expected. + for hid in hash_ids: + if args.op == "write_verify": + _write_one_verify( + svc=svc, + hash_id=int(hid), + index_prefix=index_prefix, + pages_per_block=pages_per_block, + page_bytes=int(args.page_bytes), + indexer=indexer, + kvcache=kvcache, + node_id=str(args.node_id), + redis_endpoint=redis, + etcd_endpoint=etcd, + num_shard=int(args.num_shard), + debug_dump_file=str(args.debug_dump_file or ""), + ) + else: + _ = _read_one_verify( + svc=svc, + hash_id=int(hid), + pages_per_block=pages_per_block, + page_bytes=int(args.page_bytes), + indexer=indexer, + kvcache=kvcache, + ) + + elif args.op == "loop_rw_verify": + # Continuous RW with correctness verification. + # Design goals: + # - Always verify data when key is expected resident. + # - Force LRU by writing a long stream of unique keys. + # - Probe older keys to confirm eviction happens at least once. + if not hash_ids: + raise ValueError("--hash-id or --hash-ids required") + base = int(hash_ids[0]) + + # For read-only benchmark, shard ownership is not required. + # Local indices are rebuilt from disk on startup, and Redis is optional. + + block_size = int(svc._c.block_size()) + pages_per_block = block_size // int(args.page_bytes) + indexer = torch.arange(pages_per_block, dtype=torch.int32) + + deadline = time.time() + float(args.duration_sec) + report_every = max(0.2, float(args.report_interval_sec)) + next_report = time.time() + report_every + + resident = deque(maxlen=max(1, int(args.resident_window))) + hot_ids = deque(maxlen=max(1, int(args.hot_set_size))) + writes = 0 + writes_hot = 0 + reads_ok = 0 + reads_miss = 0 + last_id = base - 1 + eviction_observed = False + ready_touched = False + + probe_reads_ok = 0 + probe_reads_fail = 0 + + probe_publish_every = max(0.2, float(args.probe_publish_interval_sec)) + probe_read_every = max(0.2, float(args.probe_read_interval_sec)) + next_probe_publish = time.time() + probe_publish_every + next_probe_read = time.time() + probe_read_every + + t0 = time.time() + + def _pick_fresh_id(candidate: int) -> int: + # Dedupe is hash-based: if a hash id already exists (from previous + # runs with persistent redis+storage), writes may be skipped and + # reads will return old bytes => verification mismatch. + # Ensure we pick ids that are not yet present. + tries = 0 + hid = int(candidate) + while tries < 64 and svc.query([hid]) == [True]: + hid += 1 + tries += 1 + if tries >= 64 and svc.query([hid]) == [True]: + # Too many collisions: jump to a random 63-bit id. + # Some CPP/Redis paths parse ids via stoll (signed 64-bit). + hid = random.getrandbits(63) + # Best-effort: avoid immediate collision. + for _ in range(8): + if svc.query([hid]) != [True]: + break + hid = random.getrandbits(63) + return int(hid) + + while time.time() < deadline: + hid = _pick_fresh_id(base + writes) + _write_one_verify( + svc=svc, + hash_id=int(hid), + index_prefix=index_prefix, + pages_per_block=pages_per_block, + page_bytes=int(args.page_bytes), + indexer=indexer, + kvcache=kvcache, + node_id=str(args.node_id), + redis_endpoint=redis, + etcd_endpoint=etcd, + num_shard=int(args.num_shard), + debug_dump_file=str(args.debug_dump_file or ""), + ) + resident.append(int(hid)) + hot_ids.append(int(hid)) + writes += 1 + last_id = int(hid) + + # Keep hot keys refreshed to avoid eviction flakiness and to enable + # other nodes to strongly verify cross-node reads. + # Refresh one hot key per iteration (after hot set has filled). + if len(hot_ids) >= hot_ids.maxlen: + refresh_id = int(hot_ids[0]) + _write_one_verify( + svc=svc, + hash_id=refresh_id, + index_prefix=index_prefix, + pages_per_block=pages_per_block, + page_bytes=int(args.page_bytes), + indexer=indexer, + kvcache=kvcache, + node_id=str(args.node_id), + redis_endpoint=redis, + etcd_endpoint=etcd, + num_shard=int(args.num_shard), + debug_dump_file=str(args.debug_dump_file or ""), + ) + writes_hot += 1 + # rotate + hot_ids.rotate(-1) + + # Mark readiness after first verified rw. + if not ready_touched and args.ready_file: + _touch(args.ready_file) + ready_touched = True + + # Verify some recent keys are still readable. + # Check the newest and one older entry when available. + if resident: + for chk in (resident[-1], resident[0]): + ok = _read_one_verify( + svc=svc, + hash_id=int(chk), + pages_per_block=pages_per_block, + page_bytes=int(args.page_bytes), + indexer=indexer, + kvcache=kvcache, + ) + if ok: + reads_ok += 1 + else: + # If even a recent key is missing, that indicates + # either incorrect LRU behavior or cross-node overwrite. + raise AssertionError(f"recent key missing unexpectedly: {chk}") + + # Probe an old key to ensure eviction happens eventually. + gap = int(args.evict_probe_gap) + if gap > 0 and writes > gap: + old = base + (writes - gap) + # Use a real read attempt to avoid false misses due to ownership-gated query(). + if not _read_one_verify( + svc=svc, + hash_id=int(old), + pages_per_block=pages_per_block, + page_bytes=int(args.page_bytes), + indexer=indexer, + kvcache=kvcache, + ): + eviction_observed = True + else: + # It might still be present depending on shard capacity; avoid forcing. + reads_miss += 1 + + now = time.time() + + # Publish hot ids for cross-node verification. + if args.probe_file and now >= next_probe_publish: + for pid in list(hot_ids): + _append_jsonl( + args.probe_file, + { + "ts": float(now), + "writer_node": str(args.node_id), + "hash_id": int(pid), + }, + ) + next_probe_publish = now + probe_publish_every + + # Read and verify probes written by other nodes. + if args.probe_file and now >= next_probe_read: + entries = _read_jsonl_tail(args.probe_file, max_bytes=256 * 1024) + # Filter: recent, different writer, integer hash_id. + cutoff = now - float(args.probe_ttl_sec) + candidates: list[int] = [] + for e in entries: + try: + ts = float(e.get("ts", 0.0)) + writer = str(e.get("writer_node", "")) + hid2 = int(e.get("hash_id")) + except Exception: + continue + if ts < cutoff: + continue + if not writer or writer == str(args.node_id): + continue + candidates.append(hid2) + + # De-dup while preserving order (prefer latest occurrences by scanning from end). + uniq: list[int] = [] + seen: set[int] = set() + for hid2 in reversed(candidates): + if hid2 in seen: + continue + seen.add(hid2) + uniq.append(hid2) + uniq.reverse() + + # Verify a bounded number. + max_n = max(0, int(args.probe_max_read)) + to_check = uniq[-max_n:] if max_n and len(uniq) > max_n else uniq + for hid2 in to_check: + # Retry briefly to tolerate mapping propagation on remote setups. + ok2 = False + deadline2 = time.time() + 2.0 + while time.time() < deadline2: + if _read_one_verify( + svc=svc, + hash_id=int(hid2), + pages_per_block=pages_per_block, + page_bytes=int(args.page_bytes), + indexer=indexer, + kvcache=kvcache, + ): + ok2 = True + break + time.sleep(0.05) + if ok2: + probe_reads_ok += 1 + else: + probe_reads_fail += 1 + raise AssertionError( + f"cross-node probe read failed for hash_id={hid2} (from other node). " + f"This indicates cross-node read inconsistency or unexpected eviction." + ) + + next_probe_read = now + probe_read_every + + if now >= next_report: + true_evict, true_evict_cnt = _true_lru_eviction(svc=svc) + elapsed = max(1e-6, now - t0) + bytes_moved = (writes + writes_hot) * int(block_size) + obj = { + "node_id": str(args.node_id), + "host": _hostname(), + "pid": os.getpid(), + "op": "loop_rw_verify", + "start_time": t0, + "elapsed_sec": elapsed, + "block_size": int(block_size), + "writes": int(writes), + "writes_hot": int(writes_hot), + "reads_ok": int(reads_ok), + "reads_miss": int(reads_miss), + "last_written_id": int(last_id), + "throughput_mb_s": float(bytes_moved / elapsed / (1024 * 1024)), + "eviction_observed": bool(true_evict), + "eviction_count": int(true_evict_cnt), + "eviction_observed_inferred": bool(eviction_observed), + "hot_ids": [int(x) for x in list(hot_ids)], + "probe_reads_ok": int(probe_reads_ok), + "probe_reads_fail": int(probe_reads_fail), + } + _atomic_write_json(args.stats_file, obj) + next_report = now + report_every + + elif args.op == "write_wait_recover": + if not hash_ids: + raise ValueError("--hash-id or --hash-ids required") + + # In distributed mode we force all shards to non-writable at startup to avoid + # concurrent writers corrupting shared shard files. + # For this op (used by test_07), we must wait for coordinator ownership to + # converge; otherwise create() can be rejected and Redis global index mapping + # will never be published. + if not disable_coord and etcd: + prefix = (index_prefix or "").strip() or "lightmem" + deadline = time.time() + 90.0 + while True: + owners = _scan_etcd_shard_owners( + etcd_endpoint=str(etcd), + prefix=str(prefix), + num_shards=int(args.num_shard), + ) + owned = sum(1 for o in owners.values() if o == str(args.node_id)) + if owned >= int(args.num_shard) and len(owners) >= int(args.num_shard): + break + if time.time() >= deadline: + raise TimeoutError( + "timeout waiting for shard ownership to converge " + f"(prefix={prefix} owned={owned}/{int(args.num_shard)} total={len(owners)}/{int(args.num_shard)})" + ) + time.sleep(0.2) + + block_size = int(svc._c.block_size()) + pages_per_block = block_size // int(args.page_bytes) + indexer = torch.arange(pages_per_block, dtype=torch.int32) + for hid in hash_ids: + h128s = _hash_128s_for_one_block(hash_id=hid, pages_per_block=pages_per_block) + task = svc.create(hash_128s=h128s, kv_page_indexer=indexer, mode="w") + _wait_task(task, timeout_s=60.0) + + _touch(args.phase1_file) + _wait_file(args.trigger_file, timeout_s=600.0) + + for sid in range(int(args.num_shard)): + try: + svc._c.recover_shard_to_redis(int(sid)) + except Exception: + pass + + elif args.op == "recover": + # Best-effort: iterate all shards. + for sid in range(int(args.num_shard)): + try: + svc._c.recover_shard_to_redis(int(sid)) + except Exception: + pass + + elif args.op == "stress_write": + block_size = int(svc._c.block_size()) + pages_per_block = block_size // int(args.page_bytes) + indexer = torch.arange(pages_per_block, dtype=torch.int32) + deadline = time.time() + float(args.duration_sec) + i = 0 + while time.time() < deadline: + base = hash_ids[0] if hash_ids else 123456 + hid = int(base) + i + h128s = _hash_128s_for_one_block(hash_id=hid, pages_per_block=pages_per_block) + task = svc.create(hash_128s=h128s, kv_page_indexer=indexer, mode="w") + _wait_task(task, timeout_s=60.0) + i += 1 + + elif args.op == "stress_read": + if not hash_ids: + raise ValueError("--hash-id or --hash-ids required") + hash_id = hash_ids[0] + block_size = int(svc._c.block_size()) + pages_per_block = block_size // int(args.page_bytes) + indexer = torch.arange(pages_per_block, dtype=torch.int32) + h128s = _hash_128s_for_one_block(hash_id=hash_id, pages_per_block=pages_per_block) + deadline = time.time() + float(args.duration_sec) + while time.time() < deadline: + _ = svc.query([hash_id]) + task = svc.create(hash_128s=h128s, kv_page_indexer=indexer, mode="r") + _wait_task(task, timeout_s=60.0) + + elif args.op == "bench_write": + # High-throughput batched writes. + # This intentionally skips per-block correctness verification to maximize bandwidth. + block_size = int(svc._c.block_size()) + pages_per_block = block_size // int(args.page_bytes) + if pages_per_block <= 0: + raise ValueError(f"invalid pages_per_block={pages_per_block}") + + max_blocks = int(args.num_pages) // int(pages_per_block) + batch_blocks = min(max(1, int(args.batch_blocks)), max(1, max_blocks)) + batch_pages = batch_blocks * int(pages_per_block) + indexer = torch.arange(batch_pages, dtype=torch.int32) + + # Use a fixed kvcache backing to avoid CPU fill cost. + # Data content does not affect dedupe (hash-based); hashes are unique per block. + # Ensure tensor exists and has enough pages. + if kvcache.shape[0] < batch_pages: + raise ValueError(f"kvcache too small: have {kvcache.shape[0]} pages, need {batch_pages}") + + # Base id: prefer user-provided --hash-id when present. + base = int(hash_ids[0]) if hash_ids else random.getrandbits(128) + + if not disable_coord: + _maybe_wait_for_single_node_ownership( + etcd_endpoint=etcd, + index_prefix=index_prefix, + node_id=str(args.node_id), + num_shards=int(args.num_shard), + storage_size_bytes=int(args.storage_size), + ) + + deadline = time.time() + float(args.duration_sec) + report_every = max(0.2, float(args.report_interval_sec)) + next_report = time.time() + report_every + + writes = 0 + eviction_observed = False + ready_touched = False + + # Estimated total capacity in blocks across all shards. + # Note: storage_size is total across shards in LocalStorageEngine. + capacity_blocks = max(1, int(int(args.storage_size) // int(block_size))) + + coord_prefix = index_prefix + effective_capacity_blocks = int(capacity_blocks) + owned_shards = 0 + uniq_nodes = 0 + last_cap_refresh = 0.0 + if not disable_coord: + effective_capacity_blocks, owned_shards, uniq_nodes = _estimate_effective_capacity_blocks( + etcd_endpoint=etcd, + node_id=str(args.node_id), + prefix=coord_prefix, + num_shards=int(args.num_shard), + capacity_blocks_total=int(capacity_blocks), + ) + last_cap_refresh = time.time() + + t0 = time.time() + dummy = 1 + last_report_time = t0 + last_report_writes = 0 + + def _make_hash_128s_for_batch(start_id: int) -> list[int]: + nonlocal dummy + out: list[int] = [] + hid = int(start_id) + for _ in range(batch_blocks): + for _ in range(int(pages_per_block) - 1): + out.append(dummy) + dummy += 1 + out.append(hid) + hid += 1 + return out + + while time.time() < deadline: + start_id = base + writes + h128s = _make_hash_128s_for_batch(start_id) + task = svc.create(hash_128s=h128s, kv_page_indexer=indexer, mode="w") + _wait_task(task, timeout_s=60.0) + + writes += batch_blocks + + if not ready_touched and args.ready_file: + _touch(args.ready_file) + ready_touched = True + + # Probe eviction for an older key. + # NOTE: + # - In etcd_mode, writes are assigned to writable shards (not necessarily the same + # deterministic hash->shard mapping as single-node mode). + # - LRU eviction is per-shard, so "an older id" is not guaranteed to be evicted. + # To avoid missing eviction, probe several sufficiently-old ids once we exceed capacity. + gap = int(args.evict_probe_gap) + if not eviction_observed: + if gap > 0 and writes > gap: + old = base + (writes - gap) + if not _query_one_block_exists(svc=svc, hash_id=int(old), pages_per_block=int(pages_per_block)): + eviction_observed = True + + # Refresh capacity estimate periodically (handles dynamic join/leave). + if not disable_coord and (time.time() - last_cap_refresh) >= 2.0: + effective_capacity_blocks, owned_shards, uniq_nodes = _estimate_effective_capacity_blocks( + etcd_endpoint=etcd, + node_id=str(args.node_id), + prefix=coord_prefix, + num_shards=int(args.num_shard), + capacity_blocks_total=int(capacity_blocks), + ) + last_cap_refresh = time.time() + + if not eviction_observed and writes >= int(effective_capacity_blocks) + max(2 * batch_blocks, 256): + # Probe very old ids (likely evicted) across the full window. + # Keep probes small to avoid impacting bandwidth. + candidates: list[int] = [] + candidates.append(int(base)) + candidates.append(int(base + (int(effective_capacity_blocks) // 4))) + candidates.append(int(base + (int(effective_capacity_blocks) // 2))) + candidates.append(int(base + ((3 * int(effective_capacity_blocks)) // 4))) + candidates.append(int(base + max(0, writes - int(effective_capacity_blocks) - 1))) + # Dedup while preserving order. + seen: set[int] = set() + uniq: list[int] = [] + for x in candidates: + if x in seen: + continue + seen.add(x) + uniq.append(x) + # Query individually to keep the Python binding behavior consistent. + for x in uniq: + if not _query_one_block_exists(svc=svc, hash_id=int(x), pages_per_block=int(pages_per_block)): + eviction_observed = True + break + + now = time.time() + if now >= next_report: + true_evict, true_evict_cnt = _true_lru_eviction(svc=svc) + elapsed = max(1e-6, now - t0) + inst_elapsed = max(1e-6, now - last_report_time) + inst_writes = int(writes) - int(last_report_writes) + inst_bytes = int(inst_writes) * int(block_size) + bytes_moved = int(writes) * int(block_size) + obj = { + "node_id": str(args.node_id), + "host": _hostname(), + "pid": os.getpid(), + "op": "bench_write", + "start_time": t0, + "elapsed_sec": elapsed, + "block_size": int(block_size), + "capacity_blocks_total": int(capacity_blocks), + "capacity_blocks_effective": int(effective_capacity_blocks), + "owned_shards": int(owned_shards), + "uniq_nodes": int(uniq_nodes), + "writes": int(writes), + "writes_hot": 0, + "reads_ok": 0, + "reads_miss": 0, + "last_written_id": int(base + writes - 1), + "throughput_mb_s": float(bytes_moved / elapsed / (1024 * 1024)), + "avg_throughput_mb_s": float(bytes_moved / elapsed / (1024 * 1024)), + "inst_throughput_mb_s": float(inst_bytes / inst_elapsed / (1024 * 1024)), + "eviction_observed": bool(true_evict), + "eviction_count": int(true_evict_cnt), + "eviction_observed_inferred": bool(eviction_observed), + "hot_ids": [], + "probe_reads_ok": 0, + "probe_reads_fail": 0, + } + _atomic_write_json(args.stats_file, obj) + next_report = now + report_every + last_report_time = now + last_report_writes = int(writes) + + elif args.op == "bench_read": + # High-throughput batched reads. + # This does NOT verify bytes; it measures read bandwidth of the hot window. + block_size = int(svc._c.block_size()) + pages_per_block = block_size // int(args.page_bytes) + if pages_per_block <= 0: + raise ValueError(f"invalid pages_per_block={pages_per_block}") + + max_blocks = int(args.num_pages) // int(pages_per_block) + batch_blocks = min(max(1, int(args.batch_blocks)), max(1, max_blocks)) + batch_pages = batch_blocks * int(pages_per_block) + indexer = torch.arange(batch_pages, dtype=torch.int32) + + if kvcache.shape[0] < batch_pages: + raise ValueError(f"kvcache too small: have {kvcache.shape[0]} pages, need {batch_pages}") + + base = int(hash_ids[0]) if hash_ids else random.getrandbits(128) + + if not disable_coord: + _maybe_wait_for_single_node_ownership( + etcd_endpoint=etcd, + index_prefix=index_prefix, + node_id=str(args.node_id), + num_shards=int(args.num_shard), + storage_size_bytes=int(args.storage_size), + ) + + # Clamp read window by total capacity (in blocks). + # If window >> capacity, most queried ids are guaranteed misses, and the benchmark + # becomes dominated by query/loop overhead rather than measuring read bandwidth. + capacity_blocks = max(1, int(int(args.storage_size) // int(block_size))) + window_req = max(1, int(args.read_window_blocks)) + window = min(window_req, capacity_blocks) + + deadline = time.time() + float(args.duration_sec) + report_every = max(0.2, float(args.report_interval_sec)) + next_report = time.time() + report_every + + reads = 0 + reads_miss = 0 + ready_touched = False + t0 = time.time() + dummy = 1 + last_report_time = t0 + last_report_reads = 0 + + def _make_hash_128s_for_read_batch(offset: int) -> list[int]: + nonlocal dummy + out: list[int] = [] + for j in range(batch_blocks): + hid = base + ((offset + j) % window) + for _ in range(int(pages_per_block) - 1): + out.append(dummy) + dummy += 1 + out.append(int(hid)) + return out + + def _make_hash_128s_for_present_blocks(offset: int, present_js: list[int]) -> list[int]: + """Build per-page cumulative hash_128s for only the present blocks. + + This avoids issuing read work for blocks that query() already determined are missing. + """ + nonlocal dummy + out: list[int] = [] + for j in present_js: + hid = base + ((offset + int(j)) % window) + for _ in range(int(pages_per_block) - 1): + out.append(dummy) + dummy += 1 + out.append(int(hid)) + return out + + i = 0 + while time.time() < deadline: + # Best-effort query to avoid expensive failing reads when keys are missing. + # IMPORTANT: PyLocalCacheService.query expects per-page cumulative hash_128s, + # not per-block ids. Here we query the underlying C++ service with the exact + # per-block hash strings (32-hex) matching what create() uses. + ids = [int(base + ((i + j) % window)) for j in range(batch_blocks)] + block_hashs = [format(int(x), "032x") for x in ids] + ok = svc._c.query(block_hashs) + present = sum(1 for x in ok if x) + miss = len(ok) - present + reads_miss += miss + if present > 0: + present_js = [j for j, x in enumerate(ok) if x] + # Pack present blocks densely at the front of the batch pages. + present_pages = int(present) * int(pages_per_block) + h128s = _make_hash_128s_for_present_blocks(i, present_js) + idx = indexer[:present_pages] + try: + task = svc.create(hash_128s=h128s, kv_page_indexer=idx, mode="r") + _wait_task(task, timeout_s=60.0) + except Exception: + # Tolerate transient mapping issues. + pass + reads += present + + # Advance the read offset so we actually cover the full window. + # Without this, the benchmark repeatedly queries/reads the same first batch. + i += batch_blocks + + if not ready_touched and args.ready_file: + _touch(args.ready_file) + ready_touched = True + + now = time.time() + if now >= next_report: + elapsed = max(1e-6, now - t0) + inst_elapsed = max(1e-6, now - last_report_time) + inst_reads = int(reads) - int(last_report_reads) + inst_bytes = int(inst_reads) * int(block_size) + bytes_moved = int(reads) * int(block_size) + obj = { + "node_id": str(args.node_id), + "host": _hostname(), + "pid": os.getpid(), + "op": "bench_read", + "start_time": t0, + "elapsed_sec": elapsed, + "block_size": int(block_size), + "writes": 0, + "writes_hot": 0, + "reads_ok": int(reads), + "reads_miss": int(reads_miss), + "last_written_id": None, + "throughput_mb_s": float(bytes_moved / elapsed / (1024 * 1024)), + "avg_throughput_mb_s": float(bytes_moved / elapsed / (1024 * 1024)), + "inst_throughput_mb_s": float(inst_bytes / inst_elapsed / (1024 * 1024)), + "eviction_observed": False, + "hot_ids": [], + "probe_reads_ok": 0, + "probe_reads_fail": 0, + } + _atomic_write_json(args.stats_file, obj) + next_report = now + report_every + last_report_time = now + last_report_reads = int(reads) + + i += batch_blocks + + else: + raise ValueError(f"unknown op: {args.op}") + + _touch(args.done_file) + + if float(args.hold_sec) > 0: + time.sleep(float(args.hold_sec)) + return 0 + finally: + try: + svc.close() + except Exception: + pass + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/read.py b/tests/read.py index 2a85552..b599343 100755 --- a/tests/read.py +++ b/tests/read.py @@ -29,6 +29,7 @@ storage_size=FILE_SIZE, num_shard=32, num_worker=32, + bandwidth_log=False, ) actual_page_size = service._page_size diff --git a/tests/test_utils.py b/tests/test_utils.py old mode 100644 new mode 100755 diff --git a/tests/write.py b/tests/write.py index 1d49452..1287cf2 100755 --- a/tests/write.py +++ b/tests/write.py @@ -29,6 +29,7 @@ storage_size=FILE_SIZE, num_shard=32, num_worker=32, + bandwidth_log=False, ) print("=" * 60) From 42858f1fd70a6a47d8672b3cda078e04bb317512 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Wed, 28 Jan 2026 15:00:14 +0800 Subject: [PATCH 3/6] add error log --- src/storage/local_cache_index.cpp | 38 ++++++- src/storage/local_storage_engine_journal.cpp | 49 +++++++-- src/storage/local_storage_engine_public.cpp | 22 ++++- src/storage/local_storage_engine_recovery.cpp | 99 ++++++++++++++++++- 4 files changed, 193 insertions(+), 15 deletions(-) diff --git a/src/storage/local_cache_index.cpp b/src/storage/local_cache_index.cpp index 106dc1d..386b610 100755 --- a/src/storage/local_cache_index.cpp +++ b/src/storage/local_cache_index.cpp @@ -1,6 +1,8 @@ #include "storage/local_cache_index.h" +#include #include +#include #include #include #include @@ -226,6 +228,8 @@ bool LocalCacheIndex::saveToSnapshot(const std::string &filename) { // Shared storage: allow other nodes/users to read+write snapshots. int fd = ::open(tmp_filename.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0666); if (fd < 0) { + std::fprintf(stderr, "[light_mem warning] snapshot save: open failed (file=%s errno=%d %s)\n", filename.c_str(), + errno, std::strerror(errno)); return false; } @@ -259,6 +263,8 @@ bool LocalCacheIndex::saveToSnapshot(const std::string &filename) { }; if (!write_all(&magic, sizeof(magic)) || !write_all(&version, sizeof(version)) || !write_all(&count, sizeof(count))) { + std::fprintf(stderr, "[light_mem warning] snapshot save: write header failed (file=%s errno=%d %s)\n", + filename.c_str(), errno, std::strerror(errno)); ::close(fd); std::remove(tmp_filename.c_str()); return false; @@ -279,6 +285,8 @@ bool LocalCacheIndex::saveToSnapshot(const std::string &filename) { if (!write_all(&hash_len, sizeof(hash_len)) || !write_all(hash.data(), hash_len) || !write_all(&slot_id, sizeof(slot_id)) || !write_all(&crc, sizeof(crc))) { + std::fprintf(stderr, "[light_mem warning] snapshot save: write entry failed (file=%s errno=%d %s)\n", + filename.c_str(), errno, std::strerror(errno)); ::close(fd); std::remove(tmp_filename.c_str()); return false; @@ -289,6 +297,8 @@ bool LocalCacheIndex::saveToSnapshot(const std::string &filename) { // Ensure file contents are durable before rename. if (::fsync(fd) != 0) { + std::fprintf(stderr, "[light_mem warning] snapshot save: fsync failed (file=%s errno=%d %s)\n", filename.c_str(), + errno, std::strerror(errno)); ::close(fd); std::remove(tmp_filename.c_str()); return false; @@ -297,6 +307,8 @@ bool LocalCacheIndex::saveToSnapshot(const std::string &filename) { // Atomic replace. if (std::rename(tmp_filename.c_str(), filename.c_str()) != 0) { + std::fprintf(stderr, "[light_mem warning] snapshot save: rename failed (file=%s errno=%d %s)\n", filename.c_str(), + errno, std::strerror(errno)); std::remove(tmp_filename.c_str()); return false; } @@ -323,15 +335,26 @@ bool LocalCacheIndex::loadFromSnapshot(const std::string &filename) { int fd = ::open(filename.c_str(), O_RDONLY); if (fd < 0) { + if (errno != ENOENT) { + std::fprintf(stderr, "[light_mem warning] snapshot load: open failed (file=%s errno=%d %s)\n", filename.c_str(), + errno, std::strerror(errno)); + } return false; } + int last_errno = 0; + bool last_eof = false; auto read_all = [&](void *p, size_t n) -> bool { char *buf = static_cast(p); size_t left = n; while (left > 0) { ssize_t r = ::read(fd, buf, left); - if (r <= 0) { + if (r == 0) { + last_eof = true; + return false; + } + if (r < 0) { + last_errno = errno; return false; } buf += static_cast(r); @@ -345,10 +368,19 @@ bool LocalCacheIndex::loadFromSnapshot(const std::string &filename) { uint64_t count = 0; if (!read_all(&magic, sizeof(magic)) || !read_all(&version, sizeof(version)) || !read_all(&count, sizeof(count))) { + if (last_eof) { + std::fprintf(stderr, "[light_mem warning] snapshot load: truncated header (file=%s)\n", filename.c_str()); + } else if (last_errno != 0) { + std::fprintf(stderr, "[light_mem warning] snapshot load: read header failed (file=%s errno=%d %s)\n", + filename.c_str(), last_errno, std::strerror(last_errno)); + } ::close(fd); return false; } if (magic != SNAPSHOT_MAGIC || version != SNAPSHOT_VERSION) { + std::fprintf(stderr, + "[light_mem warning] snapshot load: bad header (file=%s magic=0x%08x version=%u expect_magic=0x%08x expect_version=%u)\n", + filename.c_str(), magic, version, SNAPSHOT_MAGIC, SNAPSHOT_VERSION); ::close(fd); return false; } @@ -356,6 +388,7 @@ bool LocalCacheIndex::loadFromSnapshot(const std::string &filename) { for (uint64_t i = 0; i < count; i++) { uint32_t hash_len = 0; if (!read_all(&hash_len, sizeof(hash_len)) || hash_len > 4096) { + std::fprintf(stderr, "[light_mem warning] snapshot load: bad hash_len (file=%s)\n", filename.c_str()); ::close(fd); return false; } @@ -363,18 +396,21 @@ bool LocalCacheIndex::loadFromSnapshot(const std::string &filename) { std::string hash; hash.resize(hash_len); if (hash_len > 0 && !read_all(hash.data(), hash_len)) { + std::fprintf(stderr, "[light_mem warning] snapshot load: read hash failed (file=%s)\n", filename.c_str()); ::close(fd); return false; } uint64_t slot_id = 0; if (!read_all(&slot_id, sizeof(slot_id))) { + std::fprintf(stderr, "[light_mem warning] snapshot load: read slot_id failed (file=%s)\n", filename.c_str()); ::close(fd); return false; } uint32_t crc = 0; if (!read_all(&crc, sizeof(crc))) { + std::fprintf(stderr, "[light_mem warning] snapshot load: read crc failed (file=%s)\n", filename.c_str()); ::close(fd); return false; } diff --git a/src/storage/local_storage_engine_journal.cpp b/src/storage/local_storage_engine_journal.cpp index eb1a4c8..2e75649 100755 --- a/src/storage/local_storage_engine_journal.cpp +++ b/src/storage/local_storage_engine_journal.cpp @@ -4,6 +4,7 @@ #include "utils/fsync_compat.h" +#include #include #include #include @@ -108,7 +109,11 @@ void LocalStorageEngine::journalWorkerLoop(size_t shard_id) { if (queue_empty) { maybeCheckpoint(shard_id); } + } catch (const std::exception &e) { + std::fprintf(stderr, "[light_mem error] journalWorkerLoop: exception (shard=%zu): %s\n", shard_id, e.what()); + ok = false; } catch (...) { + std::fprintf(stderr, "[light_mem error] journalWorkerLoop: unknown exception (shard=%zu)\n", shard_id); ok = false; } @@ -169,33 +174,45 @@ void LocalStorageEngine::appendJournalRecord(size_t shard_id, uint64_t write_off off_t end = ::lseek(meta_fds_[shard_id], 0, SEEK_END); if (end < 0) { - throw std::runtime_error("Failed to seek meta file end"); + const int err = errno; + throw std::runtime_error(std::string("Failed to seek meta file end, errno=") + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); } if (!pwriteAll(meta_fds_[shard_id], &rec, sizeof(JournalRecord), end)) { - throw std::runtime_error("Failed to append journal header"); + const int err = errno; + throw std::runtime_error(std::string("Failed to append journal header, errno=") + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); } end += static_cast(sizeof(JournalRecord)); if (!hash.empty()) { if (!pwriteAll(meta_fds_[shard_id], hash.data(), hash.size(), end)) { - throw std::runtime_error("Failed to append journal hash"); + const int err = errno; + throw std::runtime_error(std::string("Failed to append journal hash, errno=") + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); } end += static_cast(hash.size()); } if (!evicted_hash.empty()) { if (!pwriteAll(meta_fds_[shard_id], evicted_hash.data(), evicted_hash.size(), end)) { - throw std::runtime_error("Failed to append journal evicted_hash"); + const int err = errno; + throw std::runtime_error(std::string("Failed to append journal evicted_hash, errno=") + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); } end += static_cast(evicted_hash.size()); } if (!pwriteAll(meta_fds_[shard_id], &record_crc, sizeof(uint32_t), end)) { - throw std::runtime_error("Failed to append journal crc"); + const int err = errno; + throw std::runtime_error(std::string("Failed to append journal crc, errno=") + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); } if (cache::utils::fdatasync_compat(meta_fds_[shard_id]) != 0) { - throw std::runtime_error("Failed to fdatasync meta"); + const int err = errno; + throw std::runtime_error(std::string("Failed to fdatasync meta, errno=") + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); } } @@ -218,19 +235,27 @@ void LocalStorageEngine::appendEpochMarkerLocked(size_t shard_id, uint64_t epoch off_t end = ::lseek(meta_fds_[shard_id], 0, SEEK_END); if (end < 0) { - throw std::runtime_error("Failed to seek meta file end (epoch marker)"); + const int err = errno; + throw std::runtime_error(std::string("Failed to seek meta file end (epoch marker), errno=") + + std::to_string(err) + ", reason=" + std::string(::strerror(err))); } if (!pwriteAll(meta_fds_[shard_id], &rec, sizeof(JournalRecord), end)) { - throw std::runtime_error("Failed to append epoch marker"); + const int err = errno; + throw std::runtime_error(std::string("Failed to append epoch marker, errno=") + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); } end += static_cast(sizeof(JournalRecord)); if (!pwriteAll(meta_fds_[shard_id], &record_crc, sizeof(uint32_t), end)) { - throw std::runtime_error("Failed to append epoch marker crc"); + const int err = errno; + throw std::runtime_error(std::string("Failed to append epoch marker crc, errno=") + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); } if (cache::utils::fdatasync_compat(meta_fds_[shard_id]) != 0) { - throw std::runtime_error("Failed to fdatasync meta (epoch marker)"); + const int err = errno; + throw std::runtime_error(std::string("Failed to fdatasync meta (epoch marker), errno=") + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); } } @@ -296,6 +321,8 @@ void LocalStorageEngine::checkpoint(size_t shard_id) { } if (::ftruncate(meta_fds_[shard_id], META_HEADER_SIZE) != 0) { + std::fprintf(stderr, "[light_mem warning] checkpoint: ftruncate meta failed for shard %zu (errno=%d %s)\n", + shard_id, errno, std::strerror(errno)); return; } (void)::fsync(meta_fds_[shard_id]); @@ -324,6 +351,8 @@ void LocalStorageEngine::truncateJournalToHeader(size_t shard_id) { return; } if (::ftruncate(meta_fds_[shard_id], META_HEADER_SIZE) != 0) { + std::fprintf(stderr, "[light_mem warning] truncateJournalToHeader: ftruncate failed for shard %zu (errno=%d %s)\n", + shard_id, errno, std::strerror(errno)); return; } (void)::fsync(meta_fds_[shard_id]); diff --git a/src/storage/local_storage_engine_public.cpp b/src/storage/local_storage_engine_public.cpp index 145cbc4..74c556f 100755 --- a/src/storage/local_storage_engine_public.cpp +++ b/src/storage/local_storage_engine_public.cpp @@ -7,6 +7,7 @@ #include "utils/fsync_compat.h" +#include #include #include #include @@ -352,16 +353,33 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { std::unique_lock lock(*io_locks_[shard_id]); if (!pwriteAll(file_fds_[shard_id], buf, block_size_, static_cast(offset_bytes))) { - throw std::runtime_error("Failed to write data"); + const int err = errno; + throw std::runtime_error(std::string("Failed to write data, errno=") + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); } if (cache::utils::fdatasync_compat(file_fds_[shard_id]) != 0) { - throw std::runtime_error("Failed to fdatasync data"); + const int err = errno; + throw std::runtime_error(std::string("Failed to fdatasync data, errno=") + std::to_string(err) + + ", reason=" + std::string(::strerror(err))); } // Compute CRC after data is durable. data_crc = ::crc32(0, reinterpret_cast(buf), block_size_); + } catch (const std::exception &e) { + std::fprintf(stderr, + "[light_mem error] write: data I/O failed (shard=%zu hash=%s offset=%zu): %s\n", + shard_id, hash.c_str(), offset_bytes, e.what()); + caches_[shard_id]->remove(hash); + eraseLocalShardHint(hash); + shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); + if (do_global_dedupe) { + (void)redis_lock_->del(lock_key); + } + return 0; } catch (...) { + std::fprintf(stderr, "[light_mem error] write: data I/O failed (shard=%zu hash=%s offset=%zu): unknown error\n", + shard_id, hash.c_str(), offset_bytes); caches_[shard_id]->remove(hash); eraseLocalShardHint(hash); shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); diff --git a/src/storage/local_storage_engine_recovery.cpp b/src/storage/local_storage_engine_recovery.cpp index ac8866b..1a7b891 100755 --- a/src/storage/local_storage_engine_recovery.cpp +++ b/src/storage/local_storage_engine_recovery.cpp @@ -4,6 +4,7 @@ #include #include +#include #include #include #include @@ -21,6 +22,7 @@ void LocalStorageEngine::recoverShardToRedis(size_t shard_id) { } RedisClient *redis = redisForShard(shard_id); if (!redis || !redis->connect()) { + std::fprintf(stderr, "[light_mem warning] recoverShardToRedis: Redis not reachable (shard=%zu)\n", shard_id); return; } @@ -44,12 +46,20 @@ void LocalStorageEngine::recoverShardToRedis(size_t shard_id) { int fd = ::open(snap_path.c_str(), O_RDONLY); if (fd >= 0) { + bool warned_snapshot = false; + int last_errno = 0; + bool last_eof = false; auto read_all = [&](void *p, size_t n) -> bool { char *buf = static_cast(p); size_t left = n; while (left > 0) { ssize_t r = ::read(fd, buf, left); - if (r <= 0) { + if (r == 0) { + last_eof = true; + return false; + } + if (r < 0) { + last_errno = errno; return false; } buf += static_cast(r); @@ -66,19 +76,39 @@ void LocalStorageEngine::recoverShardToRedis(size_t shard_id) { for (uint64_t i = 0; i < count; i++) { uint32_t hash_len = 0; if (!read_all(&hash_len, sizeof(hash_len)) || hash_len > 4096) { + if (!warned_snapshot) { + std::fprintf(stderr, "[light_mem warning] recoverShardToRedis: snapshot truncated/corrupt (file=%s)\n", + snap_path.c_str()); + warned_snapshot = true; + } break; } std::string hash; hash.resize(hash_len); if (hash_len > 0 && !read_all(hash.data(), hash_len)) { + if (!warned_snapshot) { + std::fprintf(stderr, "[light_mem warning] recoverShardToRedis: snapshot truncated/corrupt (file=%s)\n", + snap_path.c_str()); + warned_snapshot = true; + } break; } uint64_t slot = 0; if (!read_all(&slot, sizeof(slot))) { + if (!warned_snapshot) { + std::fprintf(stderr, "[light_mem warning] recoverShardToRedis: snapshot truncated/corrupt (file=%s)\n", + snap_path.c_str()); + warned_snapshot = true; + } break; } uint32_t crc = 0; if (!read_all(&crc, sizeof(crc))) { + if (!warned_snapshot) { + std::fprintf(stderr, "[light_mem warning] recoverShardToRedis: snapshot truncated/corrupt (file=%s)\n", + snap_path.c_str()); + warned_snapshot = true; + } break; } if (slot < shard_capacity) { @@ -87,8 +117,28 @@ void LocalStorageEngine::recoverShardToRedis(size_t shard_id) { crc_present.insert(hash); } } + } else { + // Header read failed or magic/version mismatch. + if (!warned_snapshot) { + if (last_eof) { + std::fprintf(stderr, "[light_mem warning] recoverShardToRedis: snapshot header truncated (file=%s)\n", + snap_path.c_str()); + } else if (last_errno != 0) { + std::fprintf(stderr, + "[light_mem warning] recoverShardToRedis: snapshot header read failed (file=%s errno=%d %s)\n", + snap_path.c_str(), last_errno, std::strerror(last_errno)); + } else { + std::fprintf(stderr, + "[light_mem warning] recoverShardToRedis: snapshot header mismatch (file=%s magic=0x%08x version=%u)\n", + snap_path.c_str(), magic, version); + } + warned_snapshot = true; + } } ::close(fd); + } else if (errno != ENOENT) { + std::fprintf(stderr, "[light_mem warning] recoverShardToRedis: open snapshot failed (file=%s errno=%d %s)\n", + snap_path.c_str(), errno, std::strerror(errno)); } } @@ -207,6 +257,8 @@ void LocalStorageEngine::recoverShardToRedisIncremental(size_t shard_id) { } RedisClient *redis = redisForShard(shard_id); if (!redis || !redis->connect()) { + std::fprintf(stderr, "[light_mem warning] recoverShardToRedisIncremental: Redis not reachable (shard=%zu)\n", + shard_id); return; } @@ -238,7 +290,10 @@ void LocalStorageEngine::recoverShardToRedisIncremental(size_t shard_id) { } cmds.push_back({"SET", redis->shardSeqKey(shard_id), std::to_string(superblock_seq_[shard_id])}); - (void)redis->pipeline(cmds); + if (!redis->pipeline(cmds)) { + std::fprintf(stderr, "[light_mem warning] recoverShardToRedisIncremental: Redis pipeline failed (shard=%zu cmds=%zu)\n", + shard_id, cmds.size()); + } } void LocalStorageEngine::recoverShardToRedisSmart(size_t shard_id) { @@ -253,19 +308,39 @@ void LocalStorageEngine::scanWalOps(size_t shard_id, size_t shard_capacity, off_ std::vector &ops) { struct stat st{}; if (::fstat(meta_fds_[shard_id], &st) != 0) { + std::fprintf(stderr, + "[light_mem warning] wal scan: fstat failed (shard=%zu fd=%d errno=%d %s)\n", + shard_id, meta_fds_[shard_id], errno, std::strerror(errno)); return; } const off_t end = st.st_size; off_t off = start_off; uint64_t max_epoch_seen = 0; + bool warned_io = false; + bool warned_crc = false; + bool warned_bad_record = false; while (off + static_cast(sizeof(uint32_t)) <= end) { + const off_t record_start = off; uint32_t magic = 0; if (!preadAll(meta_fds_[shard_id], &magic, sizeof(uint32_t), off)) { + if (!warned_io) { + std::fprintf(stderr, + "[light_mem warning] wal scan: read magic failed (shard=%zu off=%lld start=%lld end=%lld errno=%d %s)\n", + shard_id, static_cast(off), static_cast(start_off), + static_cast(end), errno, std::strerror(errno)); + warned_io = true; + } break; } if (magic != JOURNAL_MAGIC) { + // Not necessarily an error: tail may contain junk/partial writes. Log only if this happens at the scan start. + if (off == start_off) { + std::fprintf(stderr, + "[light_mem warning] wal scan: magic mismatch at start (shard=%zu off=%lld got=0x%08x expect=0x%08x)\n", + shard_id, static_cast(off), magic, JOURNAL_MAGIC); + } break; } @@ -318,6 +393,12 @@ void LocalStorageEngine::scanWalOps(size_t shard_id, size_t shard_capacity, off_ } const uint32_t expect = compute_crc32(tmp.data(), tmp.size()); if (expect != record_crc) { + if (!warned_crc) { + std::fprintf(stderr, + "[light_mem warning] wal scan: record crc mismatch (shard=%zu off=%lld)\n", + shard_id, static_cast(record_start)); + warned_crc = true; + } parsed = true; continue; } @@ -363,6 +444,13 @@ void LocalStorageEngine::scanWalOps(size_t shard_id, size_t shard_capacity, off_ } // Not a valid record at this offset; stop scanning. + if (!warned_bad_record) { + std::fprintf(stderr, + "[light_mem warning] wal scan: invalid record, stop scanning (shard=%zu off=%lld start=%lld end=%lld)\n", + shard_id, static_cast(record_start), static_cast(start_off), + static_cast(end)); + warned_bad_record = true; + } break; } } @@ -430,6 +518,8 @@ void LocalStorageEngine::recoverShard(size_t shard_id, size_t shard_capacity) { for (const auto &op : ops) { if (!op.evicted.empty()) { if (!redis->hdel(key, op.evicted)) { + std::fprintf(stderr, "[light_mem warning] recoverShard: Redis HDEL shard index failed (shard=%zu hash=%s)\n", + shard_id, op.evicted.c_str()); replay_ok = false; break; } @@ -437,11 +527,15 @@ void LocalStorageEngine::recoverShard(size_t shard_id, size_t shard_capacity) { (void)redis->hdel(redis->globalCrcKey(), op.evicted); } if (!redis->hset(key, op.hash, std::to_string(op.slot_id))) { + std::fprintf(stderr, "[light_mem warning] recoverShard: Redis HSET shard index failed (shard=%zu hash=%s)\n", + shard_id, op.hash.c_str()); replay_ok = false; break; } if (!redis->hset(redis->globalCrcKey(), op.hash, std::to_string(op.data_crc))) { + std::fprintf(stderr, "[light_mem warning] recoverShard: Redis HSET crc failed (shard=%zu hash=%s)\n", shard_id, + op.hash.c_str()); replay_ok = false; break; } @@ -451,6 +545,7 @@ void LocalStorageEngine::recoverShard(size_t shard_id, size_t shard_capacity) { } if (replay_ok) { if (!redis->setString(redis->shardSeqKey(shard_id), std::to_string(superblock_seq_[shard_id]))) { + std::fprintf(stderr, "[light_mem warning] recoverShard: Redis SET shard seq failed (shard=%zu)\n", shard_id); replay_ok = false; } } From 1285859f4392c51a1f5e96a811eb73e4d964c434 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Thu, 29 Jan 2026 16:40:38 +0800 Subject: [PATCH 4/6] feat: remove serach all shards --- src/storage/local_cache_index.cpp | 38 ++++++++ src/storage/local_cache_index.h | 12 +++ src/storage/local_storage_engine_public.cpp | 91 ++++--------------- src/storage/local_storage_engine_recovery.cpp | 17 +--- src/storage/local_storage_engine_shard.cpp | 11 --- 5 files changed, 70 insertions(+), 99 deletions(-) diff --git a/src/storage/local_cache_index.cpp b/src/storage/local_cache_index.cpp index 386b610..7d65b01 100755 --- a/src/storage/local_cache_index.cpp +++ b/src/storage/local_cache_index.cpp @@ -27,8 +27,21 @@ LocalCacheIndex::LocalCacheIndex(size_t capacity) : capacity_(capacity) { } } +void LocalCacheIndex::set_hooks(Hook on_ready, Hook on_erase) { + std::lock_guard lock(index_lock_); + on_ready_ = std::move(on_ready); + on_erase_ = std::move(on_erase); +} + void LocalCacheIndex::reset() { std::lock_guard lock(index_lock_); + + if (on_erase_) { + for (const auto &kv : index_) { + on_erase_(kv.first); + } + } + lru_list_.clear(); index_.clear(); empty_block_list_.clear(); @@ -53,12 +66,18 @@ void LocalCacheIndex::put_ready(const std::string &hash, size_t slot_id, uint32_ it->second.writing = false; it->second.crc = crc; lru_list_.splice(lru_list_.begin(), lru_list_, it->second.lru_iterator); + if (on_ready_) { + on_ready_(hash); + } return; } // If slot is already used by someone else, evict that hash. for (auto map_it = index_.begin(); map_it != index_.end(); ++map_it) { if (map_it->second.slot_id == slot_id) { + if (on_erase_) { + on_erase_(map_it->first); + } lru_list_.erase(map_it->second.lru_iterator); index_.erase(map_it); break; @@ -78,6 +97,9 @@ void LocalCacheIndex::put_ready(const std::string &hash, size_t slot_id, uint32_ const std::string victim = lru_list_.back(); auto vit = index_.find(victim); if (vit != index_.end()) { + if (on_erase_) { + on_erase_(victim); + } size_t freed = vit->second.slot_id; lru_list_.pop_back(); index_.erase(vit); @@ -87,6 +109,10 @@ void LocalCacheIndex::put_ready(const std::string &hash, size_t slot_id, uint32_ lru_list_.push_front(hash); index_[hash] = {lru_list_.begin(), slot_id, true, false, crc}; + + if (on_ready_) { + on_ready_(hash); + } } bool LocalCacheIndex::exists(const std::string &hash) { @@ -138,6 +164,11 @@ int LocalCacheIndex::acquire_slot(const std::string &hash, size_t &slot_id, std: if (!candidate_it->second.writing) { slot_id = candidate_it->second.slot_id; evicted_hash = candidate_hash; + + if (on_erase_) { + on_erase_(candidate_hash); + } + lru_list_.erase(it); index_.erase(candidate_it); eviction_count_++; @@ -175,6 +206,10 @@ void LocalCacheIndex::mark_ready(const std::string &hash, uint32_t crc) { if (crc != 0) { it->second.crc = crc; } + + if (on_ready_) { + on_ready_(hash); + } } } @@ -182,6 +217,9 @@ void LocalCacheIndex::remove(const std::string &hash) { std::lock_guard lock(index_lock_); auto it = index_.find(hash); if (it != index_.end()) { + if (on_erase_) { + on_erase_(hash); + } size_t slot_id = it->second.slot_id; lru_list_.erase(it->second.lru_iterator); index_.erase(it); diff --git a/src/storage/local_cache_index.h b/src/storage/local_cache_index.h index 79f7423..da919f4 100755 --- a/src/storage/local_cache_index.h +++ b/src/storage/local_cache_index.h @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -17,6 +18,8 @@ namespace storage { */ class LocalCacheIndex { public: + using Hook = std::function; + /** * @brief Internal structure to store LRU list iterator and exists call count */ @@ -34,6 +37,12 @@ class LocalCacheIndex { */ explicit LocalCacheIndex(size_t capacity); + // Optional hooks: + // - on_ready: called (under index_lock_) when a hash becomes ready/readable. + // - on_erase: called (under index_lock_) right before a hash is removed/evicted. + // These are intended for maintaining auxiliary indices (e.g. hash->shard hints). + void set_hooks(Hook on_ready, Hook on_erase); + void reset(); // Insert a ready mapping (used by recovery / redis warmup). @@ -88,6 +97,9 @@ class LocalCacheIndex { std::unordered_map index_; ///< Map from hash value to IndexEntry mutable std::mutex index_lock_; ///< Mutex protecting index data structures + Hook on_ready_; + Hook on_erase_; + uint64_t eviction_count_{0}; }; diff --git a/src/storage/local_storage_engine_public.cpp b/src/storage/local_storage_engine_public.cpp index 74c556f..5fe3f8d 100755 --- a/src/storage/local_storage_engine_public.cpp +++ b/src/storage/local_storage_engine_public.cpp @@ -111,6 +111,16 @@ LocalStorageEngine::LocalStorageEngine(const std::string &filename, const size_t journal_mu_[i] = std::make_unique(); journal_cv_[i] = std::make_unique(); } + + // Make localShardHint authoritative: update hint under the same LocalCacheIndex mutex + // for any ready insertions and removals/evictions. + if (online_mode_) { + for (size_t i = 0; i < shard_; i++) { + caches_[i]->set_hooks( + [this, i](const std::string &hash) { noteLocalShardHint(hash, i); }, + [this](const std::string &hash) { eraseLocalShardHint(hash); }); + } + } createOrOpenFiles(shard_storage_size); if (!online_mode_) { recoverAllShards(shard_capacity); @@ -154,32 +164,16 @@ bool LocalStorageEngine::query(const std::string &hash) { return caches_[shard_id]->exists(hash); } - // Best-effort hot path: if we already learned hash->shard locally and we still own that shard, - // validate via shard-local index and avoid Redis. - { - auto hinted = localShardHint(hash); - if (hinted.has_value() && hinted.value() < shard_) { - const size_t shard_id = hinted.value(); - if (caches_[shard_id]->exists(hash)) { - return true; - } else { - // Stale hint (evicted locally). - eraseLocalShardHint(hash); - } - } + // Authoritative local hint: if present, the hash is readable on this node. + auto hinted = localShardHint(hash); + if (hinted.has_value() && hinted.value() < shard_) { + return true; } - // Distributed mode: - // - If Redis is configured, prefer the global index to avoid O(shards) scans. - // - If Redis is unavailable, fall back to scanning local in-memory indices. + // Distributed mode: prefer Redis global index. if (redis_lock_ && redis_lock_->connect()) { return redis_lock_->hexists(redis_lock_->globalIndexKey(), hash); } - for (size_t i = 0; i < shard_; i++) { - if (caches_[i]->exists(hash)) { - return true; - } - } return false; } @@ -207,14 +201,8 @@ std::vector LocalStorageEngine::queryMany(const std::vector & const std::string &hash = hashs[i]; auto hinted = localShardHint(hash); if (hinted.has_value() && hinted.value() < shard_) { - const size_t shard_id = hinted.value(); - if (caches_[shard_id]->exists(hash)) { - ret[i] = true; - continue; - } else { - // Stale hint (evicted locally). - eraseLocalShardHint(hash); - } + ret[i] = true; + continue; } need_redis_hash.emplace_back(hash); need_redis_idx.emplace_back(i); @@ -225,7 +213,6 @@ std::vector LocalStorageEngine::queryMany(const std::vector & } // Second pass: Redis global index (batched). - // We use HMGET to avoid per-hash RTT. if (redis_lock_ && redis_lock_->connect()) { auto vals = redis_lock_->hmget(redis_lock_->globalIndexKey(), need_redis_hash); if (vals.has_value() && vals->size() == need_redis_hash.size()) { @@ -250,20 +237,8 @@ std::vector LocalStorageEngine::queryMany(const std::vector & } return ret; } - // If Redis IO/parsing fails, fall through to local scan. - } - - // Last resort: scan local in-memory indices. - for (size_t j = 0; j < need_redis_hash.size(); j++) { - const std::string &hash = need_redis_hash[j]; - for (size_t sid = 0; sid < shard_; sid++) { - if (caches_[sid]->exists(hash)) { - ret[need_redis_idx[j]] = true; - break; - } - } + // If Redis IO/parsing fails, treat as miss (no local full-scan fallback). } - return ret; } @@ -329,11 +304,6 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { return 0; } - // Maintain local hint bounds: if we evicted something, drop its hint. - if (!evicted_hash.empty()) { - eraseLocalShardHint(evicted_hash); - } - // If LRU evicted something, delete its Redis mappings *before* we overwrite the slot. // This avoids a window where stale Redis mapping points to overwritten data. if (do_global_dedupe && !evicted_hash.empty()) { @@ -371,7 +341,6 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { "[light_mem error] write: data I/O failed (shard=%zu hash=%s offset=%zu): %s\n", shard_id, hash.c_str(), offset_bytes, e.what()); caches_[shard_id]->remove(hash); - eraseLocalShardHint(hash); shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); if (do_global_dedupe) { (void)redis_lock_->del(lock_key); @@ -381,7 +350,6 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { std::fprintf(stderr, "[light_mem error] write: data I/O failed (shard=%zu hash=%s offset=%zu): unknown error\n", shard_id, hash.c_str(), offset_bytes); caches_[shard_id]->remove(hash); - eraseLocalShardHint(hash); shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); if (do_global_dedupe) { (void)redis_lock_->del(lock_key); @@ -393,7 +361,6 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { uint64_t epoch2 = 0; if (!isShardWritable(shard_id, &epoch2) || epoch2 != epoch) { caches_[shard_id]->remove(hash); - eraseLocalShardHint(hash); shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); if (do_global_dedupe) { (void)redis_lock_->del(lock_key); @@ -426,7 +393,6 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { if (!task->success) { caches_[shard_id]->remove(hash); - eraseLocalShardHint(hash); shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); if (do_global_dedupe) { (void)redis_lock_->del(lock_key); @@ -436,7 +402,6 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { // Mark as ready only after the WAL commit finishes. caches_[shard_id]->mark_ready(hash, data_crc); - noteLocalShardHint(hash, shard_id); shard_written_bytes_[shard_id].fetch_add(static_cast(block_size_), std::memory_order_relaxed); @@ -469,7 +434,7 @@ size_t LocalStorageEngine::read(char *buf, const std::string &hash) { } crc_available = (expected_crc != 0); } else { - // 1) Local-first: try hint then (if needed) scan local indices. + // 1) Local-first: try authoritative hint. auto hinted = localShardHint(hash); if (hinted.has_value() && hinted.value() < shard_) { const size_t i = hinted.value(); @@ -483,25 +448,11 @@ size_t LocalStorageEngine::read(char *buf, const std::string &hash) { } } - if (shard_id == static_cast(-1)) { - for (size_t i = 0; i < shard_; i++) { - const size_t off = caches_[i]->get_offset(hash); - if (off != static_cast(-1)) { - shard_id = i; - slot_id = off; - local_hit = true; - noteLocalShardHint(hash, i); - break; - } - } - } - // 2) Fallback: resolve via Redis global index. if (shard_id == static_cast(-1) && redis_lock_ && redis_lock_->connect()) { auto resolved = findShardInRedis(hash, &slot_id); if (resolved.has_value()) { shard_id = *resolved; - noteLocalShardHint(hash, shard_id); redis_hit = true; } } @@ -523,7 +474,6 @@ size_t LocalStorageEngine::read(char *buf, const std::string &hash) { eraseLocalShardHint(hash); return 0; } - noteLocalShardHint(hash, shard_id); const size_t offset_bytes = slot_id * block_size_; if (file_fds_[shard_id] < 0) { return 0; @@ -583,7 +533,6 @@ size_t LocalStorageEngine::read(char *buf, const std::string &hash) { const uint32_t got_crc = compute_crc32(buf, block_size_); if (got_crc != expected_crc) { caches_[shard_id]->remove(hash); - eraseLocalShardHint(hash); return 0; } } @@ -597,11 +546,9 @@ size_t LocalStorageEngine::read(char *buf, const std::string &hash) { (void)redis_lock_->hdel(redis_lock_->shardIndexKey(shard_id), hash); (void)redis_lock_->hdel(redis_lock_->globalCrcKey(), hash); caches_[shard_id]->remove(hash); - eraseLocalShardHint(hash); return 0; } caches_[shard_id]->put_ready(hash, slot_id, expected_crc); - noteLocalShardHint(hash, shard_id); } return block_size_; } diff --git a/src/storage/local_storage_engine_recovery.cpp b/src/storage/local_storage_engine_recovery.cpp index 1a7b891..b902894 100755 --- a/src/storage/local_storage_engine_recovery.cpp +++ b/src/storage/local_storage_engine_recovery.cpp @@ -471,16 +471,7 @@ void LocalStorageEngine::recoverShard(size_t shard_id, size_t shard_capacity) { ss << filename_ << "_" << shard_id << "/index"; std::string snap_path = ss.str(); bool snapshot_loaded = caches_[shard_id]->loadFromSnapshot(snap_path); - - // If snapshot is loaded, also warm up the local hash->shard hint map so distributed-mode - // reads/queries can locate the shard in O(1) without scanning all shards or hitting Redis. - if (snapshot_loaded && online_mode_) { - std::vector hashes; - caches_[shard_id]->dump_ready(hashes); - for (const auto &hash : hashes) { - noteLocalShardHint(hash, shard_id); - } - } + (void)snapshot_loaded; // 1. Scan WAL // We rely SOLELY on Snapshot + WAL. @@ -497,14 +488,8 @@ void LocalStorageEngine::recoverShard(size_t shard_id, size_t shard_capacity) { for (const auto &op : ops) { if (!op.evicted.empty()) { caches_[shard_id]->remove(op.evicted); - if (online_mode_) { - eraseLocalShardHint(op.evicted); - } } caches_[shard_id]->put_ready(op.hash, op.slot_id, op.data_crc); - if (online_mode_) { - noteLocalShardHint(op.hash, shard_id); - } } } diff --git a/src/storage/local_storage_engine_shard.cpp b/src/storage/local_storage_engine_shard.cpp index 103aa38..195e42a 100755 --- a/src/storage/local_storage_engine_shard.cpp +++ b/src/storage/local_storage_engine_shard.cpp @@ -165,17 +165,6 @@ void LocalStorageEngine::updateShardAssignments(const std::vector &shard caches_[sid]->reset(); } shard_recovered_epoch_[sid] = 0; - for (size_t b = 0; b < kLocalHintBuckets; b++) { - std::unique_lock lk(local_hint_mu_[b]); - auto &bucket = local_hash_to_shard_[b]; - for (auto it = bucket.begin(); it != bucket.end();) { - if (it->second == sid) { - it = bucket.erase(it); - } else { - ++it; - } - } - } } // In online mode, rebuild local index only for owned/writable shards when ownership changes. From 7d30cacbe7886c7fa72cb980c023b32a74bf2f44 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Fri, 30 Jan 2026 19:35:44 +0800 Subject: [PATCH 5/6] refine2 --- python/light_mem/etcd_coordinator.py | 29 +++++---------- python/light_mem/etcd_v3_http.py | 54 +++++++++++++--------------- 2 files changed, 34 insertions(+), 49 deletions(-) diff --git a/python/light_mem/etcd_coordinator.py b/python/light_mem/etcd_coordinator.py index 8123a13..6cf3491 100755 --- a/python/light_mem/etcd_coordinator.py +++ b/python/light_mem/etcd_coordinator.py @@ -52,7 +52,7 @@ class EtcdShardCoordinator(threading.Thread): Keys (under prefix): - nodes/{node_id} (lease-bound) - - shards/{sid}/state (FREE|CLAIMED|DRAINING|SEALED) + - shards/{sid}/state (FREE|CLAIMED|DRAINING) - shards/{sid}/owner (node_id, lease-bound) - shards/{sid}/handoff_to (node_id) # request old owner to drain/release @@ -111,16 +111,11 @@ def _shards_prefix(self) -> str: def _ensure_state(self, client, sid: int) -> None: state_key = self._k(f"shards/{sid}/state") # Create default state FREE if absent - txn = client.transaction( + client.transaction( compare=[client.transactions.create(state_key) == 0], success=[client.transactions.put(state_key, "FREE")], failure=[], ) - try: - txn - except Exception: - # Best-effort; if it fails, next reconcile will retry. - return def _get_nodes(self, client) -> List[str]: prefix = self._nodes_prefix() @@ -186,7 +181,7 @@ def _reconcile_once(self, client, lease) -> None: newly_claimed: List[int] = [] for sid in range(self._num_shards): - desired = self._assignment.get(sid) or self._desired_owner(nodes, sid) + desired = self._assignment.get(sid) state_key = self._k(f"shards/{sid}/state") owner_key = self._k(f"shards/{sid}/owner") @@ -199,12 +194,6 @@ def _reconcile_once(self, client, lease) -> None: state = self._get_text(client, state_key) state = state or "FREE" - if state == "SEALED": - # never writable - self._owned_epoch.pop(sid, None) - self._draining.pop(sid, None) - continue - owner = self._get_text(client, owner_key) handoff_to = self._get_text(client, handoff_key) @@ -310,11 +299,10 @@ def _claim(self, client, sid: int, lease) -> Optional[int]: state_key = self._k(f"shards/{sid}/state") handoff_key = self._k(f"shards/{sid}/handoff_to") - # Only claim if no owner and not SEALED. + # Only claim if no owner. txn_ok, _ = client.transaction( compare=[ client.transactions.create(owner_key) == 0, - client.transactions.value(state_key) != b"SEALED", ], success=[ client.transactions.put(owner_key, self._opt.node_id, lease=lease), @@ -488,7 +476,10 @@ def _best_effort_cleanup(c) -> None: # # Without this, different nodes may operate on different shard id ranges, # leading to inconsistent ownership keys and unsafe writes. -def _ensure_cluster_num_shards(client, *, prefix: str, expected: int) -> None: +def _ensure_cluster_num_shards(endpoints: str, prefix: str, expected: int) -> None: + host, port = _parse_endpoints(endpoints) + client = EtcdV3HttpClient(host=host, port=port) + key = f"{prefix.rstrip('/')}/config/num_shards" # First node wins by creating the key. Others validate it matches. @@ -540,9 +531,7 @@ def maybe_start_etcd_coordinator( ttl = int(coord_ttl) interval = float(coord_reconcile_sec) - host, port = _parse_endpoints(endpoints) - client = EtcdV3HttpClient(host=host, port=port) - _ensure_cluster_num_shards(client, prefix=prefix, expected=int(num_shards)) + _ensure_cluster_num_shards(endpoints=endpoints, prefix=prefix, expected=int(num_shards)) opt = EtcdOptions( endpoints=endpoints, diff --git a/python/light_mem/etcd_v3_http.py b/python/light_mem/etcd_v3_http.py index 8766535..6402522 100755 --- a/python/light_mem/etcd_v3_http.py +++ b/python/light_mem/etcd_v3_http.py @@ -266,35 +266,20 @@ def lease_keepalive_once(self, lease: EtcdLease | int, *, timeout_s: Optional[fl raise RuntimeError(f"etcd lease keepalive returned invalid TTL: {msg}") return EtcdLease(id=int(lease_id), ttl=int(ttl)) - def get(self, key: str) -> Tuple[Optional[bytes], Optional[EtcdMeta]]: - k = key.encode("utf-8") - out = self._post_json("/v3/kv/range", {"key": _b64e(k), "limit": 1}) - kvs = out.get("kvs") or [] - if not kvs: - return None, None - kv = kvs[0] - try: - meta = EtcdMeta( - key=_b64d(kv.get("key", "")), - create_revision=int(kv.get("create_revision", 0) or 0), - mod_revision=int(kv.get("mod_revision", 0) or 0), - version=int(kv.get("version", 0) or 0), - ) - v = _b64d(kv.get("value", "")) if kv.get("value") is not None else None - except Exception: - return None, None - return v, meta - - def get_prefix(self, prefix: str) -> Iterator[Tuple[Optional[bytes], EtcdMeta]]: - p = prefix.encode("utf-8") - range_end = _prefix_range_end(p) - out = self._post_json( - "/v3/kv/range", - { - "key": _b64e(p), - "range_end": _b64e(range_end), - }, - ) + def _range( + self, + *, + key: bytes, + range_end: Optional[bytes] = None, + limit: Optional[int] = None, + ) -> Iterator[Tuple[Optional[bytes], EtcdMeta]]: + payload: dict = {"key": _b64e(bytes(key))} + if range_end is not None: + payload["range_end"] = _b64e(bytes(range_end)) + if limit is not None: + payload["limit"] = int(limit) + + out = self._post_json("/v3/kv/range", payload) for kv in out.get("kvs") or []: try: meta = EtcdMeta( @@ -308,6 +293,17 @@ def get_prefix(self, prefix: str) -> Iterator[Tuple[Optional[bytes], EtcdMeta]]: continue yield v, meta + def get(self, key: str) -> Tuple[Optional[bytes], Optional[EtcdMeta]]: + k = key.encode("utf-8") + for v, meta in self._range(key=k, limit=1): + return v, meta + return None, None + + def get_prefix(self, prefix: str) -> Iterator[Tuple[Optional[bytes], EtcdMeta]]: + p = prefix.encode("utf-8") + range_end = _prefix_range_end(p) + yield from self._range(key=p, range_end=range_end) + def put(self, key: str, value: str, *, lease=None) -> None: lease_id = 0 if lease is not None: From 43251bff53ef60dc756b8bc42efab2e1569e25e2 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Mon, 2 Feb 2026 23:08:51 +0800 Subject: [PATCH 6/6] fix read --- python/light_mem/etcd_coordinator.py | 4 +- src/service/local_cache_service.h | 4 +- src/storage/local_storage_engine.h | 59 ++++- src/storage/local_storage_engine_public.cpp | 258 +++++++++----------- src/storage/storage_engine.h | 75 +++--- 5 files changed, 212 insertions(+), 188 deletions(-) diff --git a/python/light_mem/etcd_coordinator.py b/python/light_mem/etcd_coordinator.py index 6cf3491..5a0098f 100755 --- a/python/light_mem/etcd_coordinator.py +++ b/python/light_mem/etcd_coordinator.py @@ -63,7 +63,7 @@ class EtcdShardCoordinator(threading.Thread): Delayed handoff: - old owner sets draining=1 locally when it sees handoff_to != self or desired owner changed - - coordinator waits until C++ reports inflight==0 before deleting owner key + - coordinator waits until C++ reports inflight==0 (writes + online local-hit reads) before deleting owner key """ def __init__(self, service, num_shards: int, opt: EtcdOptions): @@ -336,7 +336,7 @@ def _release_if_safe(self, client, sid: int, desired_owner: Optional[str]) -> No owner_key = self._k(f"shards/{sid}/owner") state_key = self._k(f"shards/{sid}/state") - # Only release if no inflight writes. + # Only release if no inflight operations (writes + online local-hit reads). inflight = int(self._svc.shard_inflight(int(sid))) if inflight != 0: return diff --git a/src/service/local_cache_service.h b/src/service/local_cache_service.h index c5066eb..49db97f 100755 --- a/src/service/local_cache_service.h +++ b/src/service/local_cache_service.h @@ -460,7 +460,6 @@ class LocalCacheService : public CacheService { // 2. Temporary failure (all slots busy, I/O error) // 3. Write was skipped // This is acceptable for cache operations - treat as success to avoid abort - if (written == block_size_) { total_written_bytes_.fetch_add(static_cast(written), std::memory_order_relaxed); } @@ -494,8 +493,7 @@ class LocalCacheService : public CacheService { const int64_t page_bytes = page_size; // Fast path: if destination pages are contiguous in memory and indices form a contiguous range, - // we can copy the entire span in one memcpy. This is the common case for benchmarks using - // kv_page_indexer=torch.arange(...). + // we can copy the entire span in one memcpy. if (num_of_page > 0 && page_stride == page_bytes) { const int32_t first = page_idx[0]; if (first < 0 || first >= total_pages) { diff --git a/src/storage/local_storage_engine.h b/src/storage/local_storage_engine.h index 83d4c8e..7cfb037 100755 --- a/src/storage/local_storage_engine.h +++ b/src/storage/local_storage_engine.h @@ -55,11 +55,9 @@ class LocalStorageEngine : public StorageEngine { const std::string &index_endpoint, const std::string &index_prefix = std::string()); ~LocalStorageEngine() override; - bool query(const std::string &hash) override; - // Batch query variant for high-throughput callers. // Returns one bool per hash (same order). In online mode this answers "readable on this node". - std::vector queryMany(const std::vector &hashs); + std::vector queryMany(const std::vector &hashs) override; size_t write(const char *buf, const std::string &hash) override; size_t read(char *buf, const std::string &hash) override; @@ -81,8 +79,11 @@ class LocalStorageEngine : public StorageEngine { // This is a best-effort optimization; correctness is still guarded by CRC checks. void recoverShardToRedisSmart(size_t shard_id); - // Observability: number of in-flight write operations targeting a shard. - // This is best-effort and intended for control-plane decisions (e.g., waiting for draining). + // Observability: number of in-flight operations targeting a shard. + // This includes: + // - write operations + // - online-mode local-hit fast-path reads (that bypass Redis+CRC) + // It is best-effort and intended for control-plane decisions (e.g., waiting for draining). uint32_t shardInflight(size_t shard_id) const; uint64_t shardWrittenBytes(size_t shard_id) const; @@ -103,9 +104,53 @@ class LocalStorageEngine : public StorageEngine { std::optional findShardInRedis(const std::string &hash, size_t *slot_id_out); size_t pickWritableShard(const std::string &hash) const; + struct InflightCounter { + struct Guard { + explicit Guard(std::atomic *ctr) : ctr_(ctr), active_(ctr != nullptr) { + if (active_) { + ctr_->fetch_add(1, std::memory_order_relaxed); + } + } + + Guard(const Guard &) = delete; + Guard &operator=(const Guard &) = delete; + + Guard(Guard &&other) noexcept : ctr_(other.ctr_), active_(other.active_) { + other.ctr_ = nullptr; + other.active_ = false; + } + Guard &operator=(Guard &&) = delete; + + ~Guard() { + if (active_ && ctr_) { + ctr_->fetch_sub(1, std::memory_order_relaxed); + } + } + + private: + std::atomic *ctr_; + bool active_; + }; + + Guard acquire() const { return Guard(&ctr_); } + uint32_t load(std::memory_order order) const { return ctr_.load(order); } + void store(uint32_t v, std::memory_order order) { ctr_.store(v, order); } + + private: + mutable std::atomic ctr_{0}; + }; + // Best-effort local hint to avoid O(shards) scans and repeated Redis lookups. // Only used in online_mode_ and always re-validated under the shard io lock. - std::optional localShardHint(const std::string &hash) const; + struct LocalShardHintHit { + size_t shard_id; + InflightCounter::Guard inflight_guard; + }; + + // Returns a shard-local hint plus an inflight guard. + // Holding the returned guard keeps shardInflight(shard_id) > 0, so shard handoff waits for + // any local-hit read that is about to proceed. + std::optional localShardHint(const std::string &hash) const; void noteLocalShardHint(const std::string &hash, size_t shard_id); void eraseLocalShardHint(const std::string &hash); @@ -179,7 +224,7 @@ class LocalStorageEngine : public StorageEngine { std::unique_ptr[]> shard_writable_; std::unique_ptr[]> shard_draining_; std::unique_ptr[]> shard_epoch_cache_; - std::unique_ptr[]> shard_inflight_; + std::unique_ptr shard_inflight_; std::unique_ptr[]> shard_written_bytes_; std::vector journal_threads_; diff --git a/src/storage/local_storage_engine_public.cpp b/src/storage/local_storage_engine_public.cpp index 5fe3f8d..6af11c1 100755 --- a/src/storage/local_storage_engine_public.cpp +++ b/src/storage/local_storage_engine_public.cpp @@ -17,17 +17,50 @@ namespace cache { namespace storage { -std::optional LocalStorageEngine::localShardHint(const std::string &hash) const { +std::optional LocalStorageEngine::localShardHint(const std::string &hash) const { + if (!online_mode_) { + std::fprintf(stderr, "[light_mem error] localShardHint called in offline mode\n"); + return std::nullopt; + } + const size_t b = localHintBucket(hash); std::shared_lock lk(local_hint_mu_[b]); auto it = local_hash_to_shard_[b].find(hash); if (it == local_hash_to_shard_[b].end()) { return std::nullopt; } - return it->second; + const size_t shard_id = it->second; + if (shard_id >= shard_) { + return std::nullopt; + } + + // Acquire inflight inside localShardHint(). This closes the window between a successful hint + // check and the caller starting the local-hit read path. + auto inflight_guard = shard_inflight_[shard_id].acquire(); + + // In online mode, a shard in draining state is treated as not eligible for local-hit fast path. + // Re-check after acquiring inflight to avoid returning a hint during/after handoff. + if (!isShardWritable(shard_id, nullptr)) { + return std::nullopt; + } + + return LocalShardHintHit{shard_id, std::move(inflight_guard)}; } void LocalStorageEngine::noteLocalShardHint(const std::string &hash, size_t shard_id) { + if (!online_mode_) { + std::fprintf(stderr, "[light_mem error] noteLocalShardHint called in offline mode\n"); + return; + } + + if (shard_id >= shard_) { + return; + } + // Only record hints for shards eligible for local-hit fast path. + if (!isShardWritable(shard_id, nullptr)) { + return; + } + // Fast-path: avoid exclusive lock when the mapping is already up-to-date. const size_t b = localHintBucket(hash); { @@ -42,6 +75,11 @@ void LocalStorageEngine::noteLocalShardHint(const std::string &hash, size_t shar } void LocalStorageEngine::eraseLocalShardHint(const std::string &hash) { + if (!online_mode_) { + std::fprintf(stderr, "[light_mem error] eraseLocalShardHint called in offline mode\n"); + return; + } + // Fast-path: avoid exclusive lock if the key is absent. const size_t b = localHintBucket(hash); { @@ -77,7 +115,7 @@ LocalStorageEngine::LocalStorageEngine(const std::string &filename, const size_t shard_writable_ = std::make_unique[]>(shard_); shard_draining_ = std::make_unique[]>(shard_); shard_epoch_cache_ = std::make_unique[]>(shard_); - shard_inflight_ = std::make_unique[]>(shard_); + shard_inflight_ = std::make_unique(shard_); shard_written_bytes_ = std::make_unique[]>(shard_); for (size_t i = 0; i < shard_; i++) { shard_writable_[i].store(1, std::memory_order_relaxed); @@ -158,25 +196,6 @@ LocalStorageEngine::~LocalStorageEngine() { cleanup(); } -bool LocalStorageEngine::query(const std::string &hash) { - if (!online_mode_) { - const size_t shard_id = getShard(hash); - return caches_[shard_id]->exists(hash); - } - - // Authoritative local hint: if present, the hash is readable on this node. - auto hinted = localShardHint(hash); - if (hinted.has_value() && hinted.value() < shard_) { - return true; - } - - // Distributed mode: prefer Redis global index. - if (redis_lock_ && redis_lock_->connect()) { - return redis_lock_->hexists(redis_lock_->globalIndexKey(), hash); - } - return false; -} - std::vector LocalStorageEngine::queryMany(const std::vector &hashs) { std::vector ret; ret.assign(hashs.size(), false); @@ -199,10 +218,13 @@ std::vector LocalStorageEngine::queryMany(const std::vector & for (size_t i = 0; i < hashs.size(); i++) { const std::string &hash = hashs[i]; - auto hinted = localShardHint(hash); - if (hinted.has_value() && hinted.value() < shard_) { - ret[i] = true; - continue; + if (auto hinted = localShardHint(hash); hinted.has_value() && hinted->shard_id < shard_) { + const size_t sid = hinted->shard_id; + if (sid < shard_ && caches_[sid] && caches_[sid]->exists(hash)) { + ret[i] = true; + continue; + } + eraseLocalShardHint(hash); } need_redis_hash.emplace_back(hash); need_redis_idx.emplace_back(i); @@ -281,6 +303,7 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { return 0; } + auto inflight_guard = shard_inflight_[shard_id].acquire(); uint64_t epoch = 0; if (!isShardWritable(shard_id, &epoch)) { if (do_global_dedupe) { @@ -289,15 +312,12 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { return 0; } - shard_inflight_[shard_id].fetch_add(1, std::memory_order_relaxed); - size_t slot_id = 0; std::string evicted_hash; // 1. Acquire slot (LRU) int result = caches_[shard_id]->acquire_slot(hash, slot_id, evicted_hash); if (result < 0) { - shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); if (do_global_dedupe) { (void)redis_lock_->del(lock_key); } @@ -341,7 +361,6 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { "[light_mem error] write: data I/O failed (shard=%zu hash=%s offset=%zu): %s\n", shard_id, hash.c_str(), offset_bytes, e.what()); caches_[shard_id]->remove(hash); - shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); if (do_global_dedupe) { (void)redis_lock_->del(lock_key); } @@ -350,7 +369,6 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { std::fprintf(stderr, "[light_mem error] write: data I/O failed (shard=%zu hash=%s offset=%zu): unknown error\n", shard_id, hash.c_str(), offset_bytes); caches_[shard_id]->remove(hash); - shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); if (do_global_dedupe) { (void)redis_lock_->del(lock_key); } @@ -361,7 +379,6 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { uint64_t epoch2 = 0; if (!isShardWritable(shard_id, &epoch2) || epoch2 != epoch) { caches_[shard_id]->remove(hash); - shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); if (do_global_dedupe) { (void)redis_lock_->del(lock_key); } @@ -393,7 +410,6 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { if (!task->success) { caches_[shard_id]->remove(hash); - shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); if (do_global_dedupe) { (void)redis_lock_->del(lock_key); } @@ -404,8 +420,6 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { caches_[shard_id]->mark_ready(hash, data_crc); shard_written_bytes_[shard_id].fetch_add(static_cast(block_size_), std::memory_order_relaxed); - - shard_inflight_[shard_id].fetch_sub(1, std::memory_order_relaxed); if (do_global_dedupe) { (void)redis_lock_->del(lock_key); } @@ -420,137 +434,103 @@ size_t LocalStorageEngine::write(const char *buf, const std::string &hash) { } size_t LocalStorageEngine::read(char *buf, const std::string &hash) { - size_t shard_id = static_cast(-1); - size_t slot_id = static_cast(-1); - bool local_hit = false; - bool redis_hit = false; - bool crc_available = false; - uint32_t expected_crc = 0; - if (!online_mode_) { - shard_id = getShard(hash); - if (!caches_[shard_id]->get_offset_and_crc(hash, slot_id, expected_crc)) { + auto shard_id = getShard(hash); + auto slot_id = caches_[shard_id]->get_offset(hash); + if (slot_id == static_cast(-1)) { return 0; } - crc_available = (expected_crc != 0); - } else { - // 1) Local-first: try authoritative hint. - auto hinted = localShardHint(hash); - if (hinted.has_value() && hinted.value() < shard_) { - const size_t i = hinted.value(); - const size_t off = caches_[i]->get_offset(hash); - if (off != static_cast(-1)) { - shard_id = i; - slot_id = off; - local_hit = true; - } else { - eraseLocalShardHint(hash); - } - } - - // 2) Fallback: resolve via Redis global index. - if (shard_id == static_cast(-1) && redis_lock_ && redis_lock_->connect()) { - auto resolved = findShardInRedis(hash, &slot_id); - if (resolved.has_value()) { - shard_id = *resolved; - redis_hit = true; - } - } - } - - if (shard_id == static_cast(-1) || slot_id == static_cast(-1)) { - return 0; - } - - std::shared_lock lock(*io_locks_[shard_id]); - - // Owner-local fast path (online/distributed mode): - // If the mapping is from our local in-memory index AND we currently own the shard, - // then no other node is allowed to overwrite slots in this shard. - // With the shard io lock held, re-check the local offset to avoid races with local eviction. - if (online_mode_ && local_hit) { - const size_t slot_local = caches_[shard_id]->get_offset(hash); - if (slot_local != slot_id) { - eraseLocalShardHint(hash); - return 0; - } - const size_t offset_bytes = slot_id * block_size_; if (file_fds_[shard_id] < 0) { return 0; } - if (!preadAll(file_fds_[shard_id], buf, block_size_, static_cast(offset_bytes))) { - std::fprintf(stderr, "[light_mem error] read: I/O error for hash %s at offset %zu\n", hash.c_str(), - offset_bytes); - return 0; + + { + std::shared_lock lock(*io_locks_[shard_id]); + const size_t slot_local = caches_[shard_id]->get_offset(hash); + if (slot_local != slot_id) { + return 0; + } + const size_t offset_bytes = slot_id * block_size_; + if (!preadAll(file_fds_[shard_id], buf, block_size_, static_cast(offset_bytes))) { + std::fprintf(stderr, "[light_mem error] read: I/O error for hash %s at offset %zu\n", hash.c_str(), + offset_bytes); + return 0; + } } return block_size_; } - // If mapping came from Redis (or Redis is available), re-check it after taking the shard lock. - // This prevents returning wrong data if the owner evicted/overwrote the slot between lookup and read. - if (redis_lock_ && redis_lock_->connect()) { - size_t slot2 = static_cast(-1); - auto shard2 = findShardInRedis(hash, &slot2); - if (!shard2.has_value() || *shard2 != shard_id || slot2 != slot_id) { - return 0; + // Online mode: + // 1) Local-hit path (authoritative local hint). Behaves like offline mode: read locally and return. + // local hits do NOT require CRC validation. + if (auto hinted = localShardHint(hash); hinted.has_value() && hinted->shard_id < shard_) { + const size_t i = hinted->shard_id; + std::shared_lock lock(*io_locks_[i]); + const size_t slot_id = caches_[i]->get_offset(hash); + if (slot_id != static_cast(-1)) { + if (file_fds_[i] < 0) { + return 0; + } + const size_t offset_bytes = slot_id * block_size_; + if (!preadAll(file_fds_[i], buf, block_size_, static_cast(offset_bytes))) { + std::fprintf(stderr, "[light_mem error] read: I/O error for hash %s at offset %zu\n", hash.c_str(), offset_bytes); + return 0; + } + return block_size_; } + eraseLocalShardHint(hash); } - // If Redis was used to resolve, require CRC to be present before reading data. - if (redis_hit && redis_lock_ && redis_lock_->connect()) { - auto crc_s = redis_lock_->hget(redis_lock_->globalCrcKey(), hash); - if (!crc_s.has_value() || crc_s->empty()) { - (void)redis_lock_->hdel(redis_lock_->globalIndexKey(), hash); - (void)redis_lock_->hdel(redis_lock_->shardIndexKey(shard_id), hash); - (void)redis_lock_->hdel(redis_lock_->globalCrcKey(), hash); - return 0; - } - try { - expected_crc = static_cast(std::stoul(*crc_s)); - crc_available = true; - } catch (...) { - (void)redis_lock_->hdel(redis_lock_->globalIndexKey(), hash); - (void)redis_lock_->hdel(redis_lock_->shardIndexKey(shard_id), hash); - (void)redis_lock_->hdel(redis_lock_->globalCrcKey(), hash); - return 0; + // 2) Redis-resolved path (with CRC validation). + size_t shard_id = static_cast(-1); + size_t slot_id = static_cast(-1); + uint32_t expected_crc = 0; + + if (redis_lock_ && redis_lock_->connect()) { + auto resolved = findShardInRedis(hash, &slot_id); + if (resolved.has_value()) { + shard_id = *resolved; } } - const size_t offset_bytes = slot_id * block_size_; - if (file_fds_[shard_id] < 0) { + if (shard_id == static_cast(-1) || slot_id == static_cast(-1)) { return 0; } - if (!preadAll(file_fds_[shard_id], buf, block_size_, static_cast(offset_bytes))) { - std::fprintf(stderr, "[light_mem error] read: I/O error for hash %s at offset %zu\n", hash.c_str(), offset_bytes); + if (file_fds_[shard_id] < 0 || shard_id >= shard_) { return 0; } - // Local (offline) correctness: guard against a narrow TOCTOU window where a slot can be evicted - // and reused between index lookup and the actual I/O. In offline mode we don't have Redis to - // re-validate mapping, so use the locally recorded CRC (written at commit time) as a cheap - // correctness check. - if (!online_mode_ && crc_available) { - const uint32_t got_crc = compute_crc32(buf, block_size_); - if (got_crc != expected_crc) { - caches_[shard_id]->remove(hash); - return 0; + // Online mode correctness for Redis-resolved reads (or reads while shard is not writable / draining): verify CRC. + bool crc_available = false; + if (redis_lock_ && redis_lock_->connect()) { + auto crc_s = redis_lock_->hget(redis_lock_->globalCrcKey(), hash); + if (crc_s.has_value() && !crc_s->empty()) { + try { + expected_crc = static_cast(std::stoul(*crc_s)); + crc_available = (expected_crc != 0); + } catch (...) { + crc_available = false; + } } } - // CRC verification for Redis-resolved reads. - if (redis_hit && crc_available) { - const uint32_t got_crc = compute_crc32(buf, block_size_); - if (got_crc != expected_crc) { - // Stale Redis mapping: delete best-effort and return miss. - (void)redis_lock_->hdel(redis_lock_->globalIndexKey(), hash); - (void)redis_lock_->hdel(redis_lock_->shardIndexKey(shard_id), hash); - (void)redis_lock_->hdel(redis_lock_->globalCrcKey(), hash); - caches_[shard_id]->remove(hash); + if (crc_available) { + const size_t offset_bytes = slot_id * block_size_; + if (!preadAll(file_fds_[shard_id], buf, block_size_, static_cast(offset_bytes))) { + std::fprintf(stderr, "[light_mem error] read: I/O error for hash %s at offset %zu\n", hash.c_str(), offset_bytes); return 0; } - caches_[shard_id]->put_ready(hash, slot_id, expected_crc); + const uint32_t got_crc = compute_crc32(buf, block_size_); + if (got_crc == expected_crc) { + return block_size_; + } } - return block_size_; + if (redis_lock_ && redis_lock_->connect()) { + (void)redis_lock_->hdel(redis_lock_->globalIndexKey(), hash); + (void)redis_lock_->hdel(redis_lock_->shardIndexKey(shard_id), hash); + (void)redis_lock_->hdel(redis_lock_->globalCrcKey(), hash); + } + return 0; } std::shared_ptr LocalStorageEngine::getHashInfo() { diff --git a/src/storage/storage_engine.h b/src/storage/storage_engine.h index 27d9634..00f6007 100755 --- a/src/storage/storage_engine.h +++ b/src/storage/storage_engine.h @@ -1,37 +1,38 @@ -#pragma once - -#include "core/cache_task.h" -#include "core/error.h" - -#include - -namespace cache::storage { - -/** - * @brief 存储引擎基类 - */ -class StorageEngine { -public: - virtual ~StorageEngine() = default; - - /** - * @brief 查询存储引擎是否保存了给定的哈希值 - */ - virtual bool query(const std::string &hash) = 0; - - /** - * @brief 写入给定的数据到存储引擎 - */ - virtual size_t write(const char *buf, const std::string &hash) = 0; - - /** - * @brief 从存储引擎中读取数据，如果 hash 值不存在，返回0，否则返回读取的字节数 - */ - virtual size_t read(char *buf, const std::string &hash) = 0; - -protected: - // TODO 存储分页 - std::mutex lock_; // 存储锁 -}; - -} // namespace cache::storage +#pragma once + +#include "core/cache_task.h" +#include "core/error.h" + +#include +#include + +namespace cache::storage { + +/** + * @brief 存储引擎基类 + */ +class StorageEngine { +public: + virtual ~StorageEngine() = default; + + /** + * @brief 批量查询存储引擎是否保存了给定的哈希值列表 + */ + virtual std::vector queryMany(const std::vector &hashs) = 0; + + /** + * @brief 写入给定的数据到存储引擎 + */ + virtual size_t write(const char *buf, const std::string &hash) = 0; + + /** + * @brief 从存储引擎中读取数据，如果 hash 值不存在，返回0，否则返回读取的字节数 + */ + virtual size_t read(char *buf, const std::string &hash) = 0; + +protected: + // TODO 存储分页 + std::mutex lock_; // 存储锁 +}; + +} // namespace cache::storage