From 03673eaef62f8e673d926cc6d53887f1359bdc80 Mon Sep 17 00:00:00 2001 From: ykiko Date: Sat, 18 Jul 2026 04:17:34 +0800 Subject: [PATCH] fix(server): bound loaded preamble states with an LRU MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pch_cache's PreambleState blobs (mmap of the paired index, tens of MB per distinct preamble key on real projects) were opened on first use and retained for the server's lifetime: didClose, idle time and store eviction all failed to release them, so memory grew with every preamble key ever touched. - Workspace keeps a loaded-state LRU, budgeted at open documents + 2 (provider wired by the master; the slack keeps a just-closed state warm for quick reopen). Loads and rebuilds touch the LRU; didClose shrinks the budget and enforces it. Unloading only drops the cache's reference — consumers hold shared_ptr copies and finish safely, and the next use reopens the blob from disk. - CacheStore records blobs its LRU evicts (copied before the entry erase — ASan caught a use-after-free in an earlier draft) and the master drains the records on its checkpoint task, dropping orphaned pch_cache metadata only when the store still lacks the blob, so a key rebuilt after the eviction keeps its live entry. Tests: loaded states converge to the budget after closing documents and queries survive the unload/reload cycle; identical preambles keep sharing one loaded state; the store's eviction reporting is pinned drain-once at unit level. --- src/server/compiler/compiler.cpp | 2 + src/server/state/workspace.cpp | 51 ++++++++++++- src/server/state/workspace.h | 26 ++++++- src/server/transport/master_server.cpp | 32 ++++++++ src/server/transport/master_server.h | 4 + src/support/cache_store.cpp | 12 +++ src/support/cache_store.h | 14 ++++ .../server/test_memory_ownership.py | 74 ++++++++++++++++++- tests/unit/support/cache_store_tests.cpp | 9 +++ 9 files changed, 220 insertions(+), 4 deletions(-) diff --git a/src/server/compiler/compiler.cpp b/src/server/compiler/compiler.cpp index 7a1bdf3ce..9c2102514 100644 --- a/src/server/compiler/compiler.cpp +++ b/src/server/compiler/compiler.cpp @@ -671,6 +671,8 @@ kota::task Compiler::ensure_pch(Session& session, // Replace the previous blob's mapping (same key, rebuilt content); // in-flight holders of the old shared_ptr stay valid. st.state = committed.value().state; + workspace.touch_loaded_state(pch_key); + workspace.enforce_loaded_budget(); LOG_INFO("PCH built for {}: {}", path, st.path); diff --git a/src/server/state/workspace.cpp b/src/server/state/workspace.cpp index 4f543d07b..a7026c6c9 100644 --- a/src/server/state/workspace.cpp +++ b/src/server/state/workspace.cpp @@ -176,8 +176,10 @@ void Workspace::on_file_closed(std::uint32_t path_id) { if(compile_graph && compile_graph->has_unit(path_id)) { compile_graph->update(path_id); } - // PCH entries are content-keyed and may be shared with other sessions; - // blob eviction is the CacheStore's job, so nothing to clean up here. + // PCH entries are content-keyed and may be shared with other sessions, + // so nothing entry-level to clean up — but the loaded-state budget + // shrinks with the open count, and this is the moment it does. + enforce_loaded_budget(); } std::string discover_compile_commands(const Config& config, llvm::StringRef workspace_root) { @@ -398,9 +400,54 @@ std::shared_ptr Workspace::preamble_state(llvm::StringRef LOG_WARN("Retracting PCH pair {} with unreadable PreambleState blob", pch_key); store->invalidate("pch", pch_key); } + if(state) { + touch_loaded_state(pch_key); + enforce_loaded_budget(); + } return state; } +void Workspace::touch_loaded_state(llvm::StringRef pch_key) { + auto it = std::ranges::find(loaded_state_lru, pch_key); + if(it != loaded_state_lru.end()) { + loaded_state_lru.erase(it); + } + loaded_state_lru.insert(loaded_state_lru.begin(), pch_key.str()); +} + +void Workspace::enforce_loaded_budget() { + // Two extra slots over the open-document count: a closed file's + // recently used state survives a quick close/reopen, and a shared key + // serving several documents stays warm while its consumers churn. + // Open documents' keys always fit the budget, so an unload can only + // hit keys past the working set; the reload an unlucky consumer then + // pays (mmap + verification, on the event loop) is the accepted cost + // of bounding tens of MB per key. + // Unwired (tests, tools) assumes a small editor-like working set. + constexpr std::size_t default_open_documents = 6; + std::size_t budget = 2 + (open_documents ? open_documents() : default_open_documents); + + std::size_t kept = 0; + std::size_t i = 0; + while(i < loaded_state_lru.size()) { + auto it = pch_cache.find(loaded_state_lru[i]); + // Erased entries and already-unloaded keys just fall out of the + // list (invalidation and store eviction bypass the LRU). + if(it == pch_cache.end() || !it->second.state) { + loaded_state_lru.erase(loaded_state_lru.begin() + i); + continue; + } + if(kept < budget) { + kept += 1; + i += 1; + continue; + } + LOG_DEBUG("Unloading PreambleState of {} (budget {})", loaded_state_lru[i], budget); + it->second.state.reset(); + loaded_state_lru.erase(loaded_state_lru.begin() + i); + } +} + void Workspace::load_cache(ContextResolver& contexts) { if(!store) return; diff --git a/src/server/state/workspace.h b/src/server/state/workspace.h index 0bb340360..0cb66a392 100644 --- a/src/server/state/workspace.h +++ b/src/server/state/workspace.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include #include @@ -227,6 +228,15 @@ struct Workspace { /// of CacheStore state; blob paths come from the store. llvm::StringMap pch_cache; + /// Keys of pch_cache entries whose PreambleState is currently loaded, + /// most recently used first (see enforce_loaded_budget). + llvm::SmallVector loaded_state_lru; + + /// Open-document count provider, wired by the master. Sizes the + /// loaded-state budget; unset (tests, tools) falls back to + /// default_open_documents. + std::function open_documents; + /// Crash budget for shared build artifacts (PCH/PCM), keyed by the /// same content-derived cache keys: an artifact that keeps killing /// workers is refused until its content — and therefore its key — @@ -295,9 +305,23 @@ struct Workspace { /// on-disk pair is retracted from the store as well — otherwise every /// later session re-adopts the corrupt pair from cache.json and /// silently degrades again. With the pair gone the next ensure_pch is - /// a miss and rebuilds both halves. + /// a miss and rebuilds both halves. Loads count against the + /// loaded-state budget (see enforce_loaded_budget). std::shared_ptr preamble_state(llvm::StringRef pch_key); + /// Move a pch key to the front of the loaded-state LRU. Called + /// whenever an entry's PreambleState is opened or replaced. + void touch_loaded_state(llvm::StringRef pch_key); + + /// Unload PreambleState blobs beyond the budget (open documents + 2), + /// least recently used first. Without this every preamble key ever + /// touched keeps its blob mapped for the server's lifetime — tens of + /// MB per key on real projects, released by neither didClose nor + /// store eviction. Unloading only drops the entry's reference: + /// consumers holding the shared_ptr finish safely, and the next use + /// reopens the blob from disk. + void enforce_loaded_budget(); + /// Load PCH/PCM validity metadata plus the context resolver's slices /// from cache.json (under the store's versioned root); entries whose /// blob is gone from the store are dropped. diff --git a/src/server/transport/master_server.cpp b/src/server/transport/master_server.cpp index 241098de6..8f3792d76 100644 --- a/src/server/transport/master_server.cpp +++ b/src/server/transport/master_server.cpp @@ -43,6 +43,12 @@ MasterServer::MasterServer(kota::event_loop& loop, std::string self_path) : // lifetime and turns reports into state (notify_log) plus a wake-up // signal. Master-side reports only ever fire on the event-loop thread // (see support/anomaly.h), so no synchronization is needed here. + // The loaded-state budget follows the open-document count; Workspace + // cannot see SessionStore, so the master wires the provider. + workspace.open_documents = [this] { + return sessions.sessions.size(); + }; + logging::set_notify_hook([this](logging::NotifyLevel level, std::string_view message) { notify_log.push_back(NotifyMessage{level, std::string(message)}); if(notify_log.size() > notify_log_limit) { @@ -392,6 +398,32 @@ kota::task<> MasterServer::cache_checkpoint_task() { if(workspace.store) { // Offload to the thread pool: checkpoint writes the manifest. co_await kota::queue([this] { workspace.store->checkpoint(); }); + drain_store_evictions(); + } + } +} + +void MasterServer::drain_store_evictions() { + // The store's LRU evicted these blobs from disk (inside commit, maybe + // on a worker thread); drop the derived pch_cache metadata here on the + // event loop, or the content-keyed map grows for the server's + // lifetime even for keys never requested again. An entry mid-rebuild + // keeps its slot — its commit republishes fresh blobs over the + // eviction. + for(auto& evicted: workspace.store->take_evictions()) { + if(evicted.ns != "pch") { + continue; + } + // A key rebuilt after the eviction was recorded has live blobs + // again — the record is stale, not the entry. Erase only when the + // store still lacks the blob, and never mid-rebuild (the commit + // republishes over the eviction). + if(workspace.store->lookup("pch", evicted.key)) { + continue; + } + if(auto it = workspace.pch_cache.find(evicted.key); + it != workspace.pch_cache.end() && !it->second.building) { + workspace.pch_cache.erase(it); } } } diff --git a/src/server/transport/master_server.h b/src/server/transport/master_server.h index af7813817..2aa7064bd 100644 --- a/src/server/transport/master_server.h +++ b/src/server/transport/master_server.h @@ -207,6 +207,10 @@ class MasterServer { /// times survive crashes (the store itself is passive by design). kota::task<> cache_checkpoint_task(); + /// Drop pch_cache metadata for blobs the store's LRU evicted from + /// disk (see cache_checkpoint_task). + void drain_store_evictions(); + /// The file tracker's polling loops: each tick hands the tracker's /// event batch to dispatch(). Spawned by initialize() when the /// configured interval is non-zero. diff --git a/src/support/cache_store.cpp b/src/support/cache_store.cpp index a9889a316..788a41569 100644 --- a/src/support/cache_store.cpp +++ b/src/support/cache_store.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #ifdef _WIN32 @@ -224,6 +225,9 @@ struct CacheStore::State { void evict_locked(Namespace& ns, llvm::StringRef keep_key); void checkpoint_locked(); + + /// Evictions recorded for take_evictions(); guarded by mutex. + std::vector evictions; }; CacheStore::CacheStore(std::unique_ptr state) : state(std::move(state)) {} @@ -640,6 +644,11 @@ void CacheStore::invalidate(llvm::StringRef ns, llvm::StringRef key) { maybe_checkpoint(); } +std::vector CacheStore::take_evictions() { + std::lock_guard guard(state->mutex); + return std::exchange(state->evictions, {}); +} + std::size_t CacheStore::pending_tmp_files() const { std::lock_guard guard(state->mutex); std::size_t count = 0; @@ -731,6 +740,9 @@ void CacheStore::State::evict_locked(Namespace& ns, llvm::StringRef keep_key) { candidate.size); auto it = ns.entries.find(candidate.key); ns.total_size -= it->second.size; + // Record before the erase: the candidate's key view borrows the + // entry's own storage. + evictions.push_back({ns.config.name, candidate.key.str()}); ns.entries.erase(it); dirty = true; } diff --git a/src/support/cache_store.h b/src/support/cache_store.h index dc6fcc480..42a4025ad 100644 --- a/src/support/cache_store.h +++ b/src/support/cache_store.h @@ -6,6 +6,7 @@ #include #include #include +#include #include "llvm/ADT/STLFunctionalExtras.h" #include "llvm/ADT/StringRef.h" @@ -195,6 +196,19 @@ class CacheStore { /// for the stats endpoint; scans the directory, so not free. std::size_t pending_tmp_files() const; + /// A blob the LRU budget evicted, reported so owners of derived + /// in-memory state (e.g. the PCH cache's entry metadata) can drop it: + /// eviction happens inside commit — possibly on a worker thread — so + /// the store records instead of calling back, and the owner drains on + /// its own loop. + struct EvictedBlob { + std::string ns; + std::string key; + }; + + /// Drain the evictions recorded since the last call. + std::vector take_evictions(); + /// The versioned root directory, e.g. `{root}/cache/v1`. Callers may /// place their own metadata files directly under it (the store only /// manages namespace subdirectories); they die with the version. diff --git a/tests/integration/server/test_memory_ownership.py b/tests/integration/server/test_memory_ownership.py index 003e43bca..fba5a99a3 100644 --- a/tests/integration/server/test_memory_ownership.py +++ b/tests/integration/server/test_memory_ownership.py @@ -10,6 +10,7 @@ from lsprotocol.types import ( DidChangeTextDocumentParams, + DidCloseTextDocumentParams, DidSaveTextDocumentParams, TextDocumentContentChangeWholeDocument, TextDocumentIdentifier, @@ -23,7 +24,12 @@ wait_for_recompile, ) from tests.tools.compile_commands import write_cdb -from tests.tools.workspace import get_field, list_tmp_files, pin_cache_to_workspace +from tests.tools.workspace import ( + doc, + get_field, + list_tmp_files, + pin_cache_to_workspace, +) async def wait_stats(client, predicate, *, timeout=30.0, message=""): @@ -162,3 +168,69 @@ async def test_cancel_storm_leaves_no_tmp(client, tmp_path): assert get_field(stats, "pchCacheEntries") >= 1, f"a PCH was built: {stats}" assert get_field(stats, "sessions") == 1, f"one open document: {stats}" assert_no_anomaly(client, tmp_path) + + +async def test_preamble_state_released(client, tmp_path): + pin_cache_to_workspace(tmp_path) + names = [] + for i in range(3): + (tmp_path / f"h{i}.h").write_text(f"#pragma once\nint distinct_{i} = {i};\n") + (tmp_path / f"m{i}.cpp").write_text( + f'#include "h{i}.h"\nint use_{i}() {{ return distinct_{i}; }}\n' + ) + names.append(f"m{i}.cpp") + write_cdb(tmp_path, names) + await client.initialize(tmp_path) + + uris = [] + for i in range(3): + uri, _ = await client.open_and_wait(tmp_path / f"m{i}.cpp") + uris.append(uri) + stats = await client.stats() + assert get_field(stats, "pchLoadedStates") == 3, ( + f"three distinct preambles: {stats}" + ) + + for uri in uris: + client.text_document_did_close( + DidCloseTextDocumentParams(text_document=doc(uri)) + ) + # Budget follows the open count (open + 2, the slack keeping a + # just-closed state warm): with everything closed at least one of the + # three states must unload instead of staying mapped forever. + await wait_stats( + client, + lambda s: get_field(s, "pchLoadedStates") <= 2, + message="closing documents must release loaded preamble states", + ) + + # Reload after unload: reopening must reopen the blob from disk and + # keep serving queries against the preamble's symbols. + uri0, _ = await client.open_and_wait(tmp_path / "m0.cpp") + hover = await client.hover_at(uri0, 1, 25) + assert hover is not None, "query must survive an unload/reload cycle" + stats = await client.stats() + assert get_field(stats, "pchLoadedStates") >= 1, f"state reloaded: {stats}" + assert_no_anomaly(client, tmp_path) + + +async def test_same_preamble_shared(client, tmp_path): + pin_cache_to_workspace(tmp_path) + (tmp_path / "shared.h").write_text("#pragma once\nint shared_val = 1;\n") + names = [] + for i in range(4): + (tmp_path / f"s{i}.cpp").write_text( + f'#include "shared.h"\nint fn_{i}() {{ return shared_val; }}\n' + ) + names.append(f"s{i}.cpp") + write_cdb(tmp_path, names) + await client.initialize(tmp_path) + + for i in range(4): + await client.open_and_wait(tmp_path / f"s{i}.cpp") + # Identical preambles share one content key, and sharing means one + # blob: opening more consumers must not multiply loaded states. + stats = await client.stats() + assert get_field(stats, "pchLoadedStates") == 1, f"one shared key: {stats}" + assert get_field(stats, "pchCacheEntries") == 1, f"one shared entry: {stats}" + assert_no_anomaly(client, tmp_path) diff --git a/tests/unit/support/cache_store_tests.cpp b/tests/unit/support/cache_store_tests.cpp index 663f860f4..60f9eb25b 100644 --- a/tests/unit/support/cache_store_tests.cpp +++ b/tests/unit/support/cache_store_tests.cpp @@ -218,6 +218,15 @@ TEST_CASE(LruEviction) { ASSERT_TRUE(store.lookup("pch", "a").has_value()); ASSERT_FALSE(store.lookup("pch", "b").has_value()); ASSERT_TRUE(store.lookup("pch", "c").has_value()); + + // The eviction is reported exactly once: owners of derived in-memory + // state drain the record on their own loop, and a second drain must + // not replay it. + auto evicted = store.take_evictions(); + ASSERT_EQ(evicted.size(), 1U); + ASSERT_EQ(evicted[0].ns, std::string("pch")); + ASSERT_EQ(evicted[0].key, std::string("b")); + ASSERT_TRUE(store.take_evictions().empty()); } TEST_CASE(FreshCommitNotEvicted) {