diff --git a/metagraph/src/cli/config/config.cpp b/metagraph/src/cli/config/config.cpp index 68726f9fde..6c7ccbffc2 100644 --- a/metagraph/src/cli/config/config.cpp +++ b/metagraph/src/cli/config/config.cpp @@ -407,8 +407,8 @@ Config::Config(int argc, char *argv[]) { relax_arity_brwt = atoi(get_value(i++)); } else if (!strcmp(argv[i], "--RA-ivbuff-size")) { RA_ivbuffer_size = atoll(get_value(i++)); - // } else if (!strcmp(argv[i], "--cache-size")) { - // row_cache_size = atoi(get_value(i++)); + } else if (!strcmp(argv[i], "--cache-size")) { + server_cache_size = std::stod(get_value(i++)); } else if (!strcmp(argv[i], "-h") || !strcmp(argv[i], "--help")) { print_welcome_message(); print_usage(argv[0], identity); @@ -1435,7 +1435,7 @@ if (advanced) { fprintf(stderr, "\t-p --parallel [INT] \tmaximum number of parallel connections [1]\n"); fprintf(stderr, "\t --threads-each [INT] \tnumber of threads per graph [1]\n"); fprintf(stderr, "\t --one-pass-brwt \tuse one-pass parallel BRWT traversal for queries [off]\n"); - // fprintf(stderr, "\t --cache-size [INT] \tnumber of uncompressed rows to store in the cache [0]\n"); + fprintf(stderr, "\t --cache-size [GB] \tmax size of the /search result cache, in GB (0 disables) [1]\n"); fprintf(stderr, "\n\t --num-top-labels [INT] \tmaximum number of top labels per query by default [10'000]\n"); fprintf(stderr, "\t --no-coord-mapping \t\tquery without mapping coords to sequence headers even if the .seq index exists [off]\n"); fprintf(stderr, "\t --mem-cap-gb [FLOAT] \tmemory in GB available for the server to load graphs for queries into RAM [0]\n"); diff --git a/metagraph/src/cli/config/config.hpp b/metagraph/src/cli/config/config.hpp index 9228e6cb63..815e748ee3 100644 --- a/metagraph/src/cli/config/config.hpp +++ b/metagraph/src/cli/config/config.hpp @@ -90,6 +90,7 @@ class Config { unsigned int arity_brwt = 2; unsigned int relax_arity_brwt = 10; unsigned long long RA_ivbuffer_size = 16'384; // in B + double server_cache_size = 1.0; // server query cache, in GB (default 1 GB; 0 disables) unsigned int min_tip_size = 1; unsigned int min_unitig_median_kmer_abundance = 1; int fallback_abundance_cutoff = 1; diff --git a/metagraph/src/cli/server.cpp b/metagraph/src/cli/server.cpp index a56aae03ff..6ea0e817f2 100644 --- a/metagraph/src/cli/server.cpp +++ b/metagraph/src/cli/server.cpp @@ -17,6 +17,7 @@ #include "query.hpp" #include "align.hpp" #include "server_utils.hpp" +#include "server_cache.hpp" #include "cli/load/load_annotation.hpp" @@ -374,130 +375,234 @@ int run_server(Config *config) { std::condition_variable space_cv; std::mutex space_mutex; - // the actual server + size_t cache_size_bytes = static_cast(config->server_cache_size * 1e9); + ServerQueryCache search_cache(cache_size_bytes); + if (cache_size_bytes) { + logger->info("[Server] /search result cache enabled, max size: {} GB", + config->server_cache_size); + } else { + logger->info("[Server] /search result cache disabled (--cache-size 0)"); + } + HttpServer server; + + // ---------- /search ---------- + // + // Request lifecycle: + // 1. process_request invokes the compute lambda below. + // 2. The compute lambda derives a cache key and `acquire()`s a + // handle. On a hit it returns the cached serialized response + // directly (no JSON re-serialization). On a miss it runs the + // (potentially multi-minute) query, serializes the JSON once, + // publishes the serialized string via Handle::set_result(), + // and returns the same string. + // 3. process_request writes the body to the wire (compressing on + // the fly if Accept-Encoding requests it). + // 4. The on_sent callback runs after the async write completes: + // ec == 0 → mark_delivered() (DELIVERED sink) + // ec != 0 → mark_protected() (entry kept with priority for + // the configured retry window) + // + // The handle has to span both lambdas so the on_sent callback can + // call mark_delivered/mark_protected on the same entry. We pass it + // through a shared_ptr captured by both lambdas. server.resource["^/search"]["POST"] = [&](shared_ptr response, shared_ptr request) { size_t request_id = num_requests++; - process_request(response, request, request_id, [&](const std::string& content) { - if (!config->fnames.size() && anno_graph.wait_for(0s) != std::future_status::ready) - throw CurrentlyInitializingError(); // the index is not loaded yet, so we can't process the request - - Json::Value content_json = parse_json_string(content); - logger->info("[Server] Request {}: {}", request_id, content_json.toStyledString()); - Json::Value result; + auto handle_holder = std::make_shared(); + process_request(response, request, request_id, + [&, handle_holder](const std::string &content) -> std::string { + if (!config->fnames.size() && anno_graph.wait_for(0s) != std::future_status::ready) + throw CurrentlyInitializingError(); // graph still loading + + Json::Value content_json = parse_json_string(content); + logger->info("[Server] Request {}: {}", request_id, content_json.toStyledString()); + + // Build graph identity for the cache key: the single graph + // path in the simple case, or a sorted concatenation of + // (graph, annotator) pairs for the multi-graph server. + std::string graph_identity; + if (!config->fnames.size()) { + graph_identity = config->infbase; + } else { + auto graphs_to_query = filter_graphs_from_list(indexes, content_json, request_id); + std::ostringstream oss; + for (const auto &name : graphs_to_query) { + for (const auto &[g, a] : indexes[name]) { + oss << g << '|' << a << ';'; + } + } + graph_identity = oss.str(); + } + std::string cache_key = make_search_cache_key(content_json, *config, graph_identity); + + // Cache handshake: hit returns the cached body directly, + // miss falls through to compute + serialize + publish. + *handle_holder = search_cache.acquire(cache_key); + if (!handle_holder->is_miss()) { + auto cached = handle_holder->get(); + if (cache_size_bytes) { + logger->info("[Server] Request {}: cache HIT {:.1f} KB " + "(occupancy {:.1f}/{:.0f} MB, {} entries)", + request_id, cached->size() / 1e3, + search_cache.size_bytes() / 1e6, + cache_size_bytes / 1e6, + search_cache.entry_count()); + } + return *cached; + } - // simple case with a single graph pair - if (!config->fnames.size()) { - if (content_json.isMember("graphs")) - throw std::invalid_argument("Bad request: no support for filtering graphs on this server"); - logger->trace("Request {}: Started querying graph {}, in total graphs being queried at the moment: {}", - request_id, config->infbase, graphs_being_queried.fetch_add(1) + 1); + // Cache miss: run the search, publish the result. If the + // compute throws, propagate the exception to any duplicate + // waiters via set_exception() before re-throwing. + Json::Value result; try { - result = process_search_request(content_json, *anno_graph.get(), *config); - } catch (...) { - graphs_being_queried--; - throw; - } - graphs_being_queried--; - } else { - std::vector graphs_to_query - = filter_graphs_from_list(indexes, content_json, request_id); - std::mutex mu; - std::vector> futures; - for (const auto &name : graphs_to_query) { - for (const auto &[graph_fname, anno_fname] : indexes[name]) { - futures.push_back(graphs_pool.enqueue([&,config,graph_fname=graph_fname,anno_fname=anno_fname]() { - logger->trace("Request {}: Started querying graph {}. In total graphs being queried at the moment: {}", - request_id, graph_fname, graphs_being_queried.fetch_add(1) + 1); - size_t index_size_reserved = 0; - auto release_memory = [&]() { - if (!index_size_reserved) - return; - { - std::unique_lock lock(space_mutex); - memory_left += index_size_reserved; - } - index_size_reserved = 0; - space_cv.notify_all(); - }; - try { - std::unique_ptr index_loaded; - const AnnotatedDBG *index; - bool in_ram = content_json.isMember("in_ram") && content_json["in_ram"].asBool(); - if (in_ram && !loaded_with_mmap) - in_ram = false; // already in RAM, no need to re-load - size_t index_size = in_ram ? std::filesystem::file_size(graph_fname) - + std::filesystem::file_size(anno_fname) - : -1; - if (in_ram && index_size > memory_all) { - logger->warn("Request {}: Graph of size {} GB is too large to fit into " - "RAM (reserved memory: {} GB). It will be queried with mmap", - request_id, index_size / 1e9, memory_all / 1e9); - in_ram = false; - } - Timer timer; - if (in_ram) { - { - std::unique_lock lock(space_mutex); - space_cv.wait(lock, [&]() { - return memory_left >= index_size; - }); - memory_left -= index_size; - } - index_size_reserved = index_size; - logger->trace("Request {}: Loading graph {} of size {} GB to RAM...", - request_id, graph_fname, index_size / 1e9); - timer.reset(); - Config config_copy = *config; - config_copy.infbase = graph_fname; - config_copy.infbase_annotators = { anno_fname }; - index_loaded = initialize_annotated_dbg(config_copy); - index = index_loaded.get(); - } else { - index = graphs_cache.at({ graph_fname, anno_fname }).get(); - } - - auto json = process_search_request(content_json, *index, *config); - logger->trace("Request {}: {} graph {} {} in {} sec", - request_id, in_ram ? "Loaded and searched" : "Searched", - graph_fname, anno_fname, timer.elapsed()); - - index_loaded.reset(); - release_memory(); - - std::lock_guard lock(mu); - if (result.empty()) { - result = std::move(json); - } else { - assert(json.size() == result.size()); - for (Json::ArrayIndex i = 0; i < result.size(); ++i) { - if (result[i][SeqSearchResult::SEQ_DESCRIPTION_JSON_FIELD] - != json[i][SeqSearchResult::SEQ_DESCRIPTION_JSON_FIELD]) { - throw std::logic_error("ERROR: Results for different sequences can't be merged"); + if (!config->fnames.size()) { + if (content_json.isMember("graphs")) + throw std::invalid_argument("Bad request: no support for filtering graphs on this server"); + logger->trace("Request {}: Started querying graph {}, in total graphs being queried at the moment: {}", + request_id, config->infbase, graphs_being_queried.fetch_add(1) + 1); + try { + result = process_search_request(content_json, *anno_graph.get(), *config); + } catch (...) { + graphs_being_queried--; + throw; + } + graphs_being_queried--; + } else { + std::vector graphs_to_query + = filter_graphs_from_list(indexes, content_json, request_id); + std::mutex mu; + std::vector> futures; + for (const auto &name : graphs_to_query) { + for (const auto &[graph_fname, anno_fname] : indexes[name]) { + futures.push_back(graphs_pool.enqueue([&,config,graph_fname=graph_fname,anno_fname=anno_fname]() { + logger->trace("Request {}: Started querying graph {}. In total graphs being queried at the moment: {}", + request_id, graph_fname, graphs_being_queried.fetch_add(1) + 1); + size_t index_size_reserved = 0; + auto release_memory = [&]() { + if (!index_size_reserved) + return; + { + std::unique_lock lock(space_mutex); + memory_left += index_size_reserved; + } + index_size_reserved = 0; + space_cv.notify_all(); + }; + try { + std::unique_ptr index_loaded; + const AnnotatedDBG *index; + bool in_ram = content_json.isMember("in_ram") && content_json["in_ram"].asBool(); + if (in_ram && !loaded_with_mmap) + in_ram = false; // already in RAM, no need to re-load + size_t index_size = in_ram ? std::filesystem::file_size(graph_fname) + + std::filesystem::file_size(anno_fname) + : -1; + if (in_ram && index_size > memory_all) { + logger->warn("Request {}: Graph of size {} GB is too large to fit into " + "RAM (reserved memory: {} GB). It will be queried with mmap", + request_id, index_size / 1e9, memory_all / 1e9); + in_ram = false; } - for (auto&& value : json[i]["results"]) { - result[i]["results"].append(std::move(value)); + Timer timer; + if (in_ram) { + { + std::unique_lock lock(space_mutex); + space_cv.wait(lock, [&]() { + return memory_left >= index_size; + }); + memory_left -= index_size; + } + index_size_reserved = index_size; + logger->trace("Request {}: Loading graph {} of size {} GB to RAM...", + request_id, graph_fname, index_size / 1e9); + timer.reset(); + Config config_copy = *config; + config_copy.infbase = graph_fname; + config_copy.infbase_annotators = { anno_fname }; + index_loaded = initialize_annotated_dbg(config_copy); + index = index_loaded.get(); + } else { + index = graphs_cache.at({ graph_fname, anno_fname }).get(); + } + + auto json = process_search_request(content_json, *index, *config); + logger->trace("Request {}: {} graph {} {} in {} sec", + request_id, in_ram ? "Loaded and searched" : "Searched", + graph_fname, anno_fname, timer.elapsed()); + + index_loaded.reset(); + release_memory(); + + std::lock_guard lock(mu); + if (result.empty()) { + result = std::move(json); + } else { + assert(json.size() == result.size()); + for (Json::ArrayIndex i = 0; i < result.size(); ++i) { + if (result[i][SeqSearchResult::SEQ_DESCRIPTION_JSON_FIELD] + != json[i][SeqSearchResult::SEQ_DESCRIPTION_JSON_FIELD]) { + throw std::logic_error("ERROR: Results for different sequences can't be merged"); + } + for (auto&& value : json[i]["results"]) { + result[i]["results"].append(std::move(value)); + } + } } + } catch (...) { + release_memory(); + graphs_being_queried--; + return std::current_exception(); } - } - } catch (...) { - release_memory(); - graphs_being_queried--; - return std::current_exception(); + graphs_being_queried--; + return std::exception_ptr(); + })); } - graphs_being_queried--; - return std::exception_ptr(); - })); + } + for (auto &future : futures) { + if (auto ex = future.get()) + std::rethrow_exception(ex); + } } + } catch (...) { + handle_holder->set_exception(std::current_exception()); + throw; } - for (auto &future : futures) { - if (auto ex = future.get()) - std::rethrow_exception(ex); + + // Serialize once: the cached form *is* the response + // body. Cache hits then skip serialization entirely. + std::string serialized = Json::writeString(Json::StreamWriterBuilder(), result); + size_t entry_size = serialized.size(); + handle_holder->set_result(serialized); + if (cache_size_bytes) { + logger->info("[Server] Request {}: cache STORE {:.1f} KB " + "(occupancy {:.1f}/{:.0f} MB, {} entries)", + request_id, entry_size / 1e3, + search_cache.size_bytes() / 1e6, + cache_size_bytes / 1e6, + search_cache.entry_count()); } - } - return result; - }); + return serialized; + }, + // on_sent: fires after the async response write completes. + // Reports the delivery outcome to the cache so subsequent + // identical requests hit a DELIVERED entry (LRU normal) or + // a PROTECTED entry (priority retention for retries). + [handle_holder, request_id, &search_cache](const SimpleWeb::error_code &ec) { + if (!handle_holder || !*handle_holder) + return; // never acquired (e.g. early throw before acquire) + if (ec) { + auto ttl_min = std::chrono::duration_cast( + search_cache.protection_ttl()).count(); + logger->info("[Server] Request {}: delivery FAILED ({}); " + "cache entry protected for {} min retry window", + request_id, ec.message(), ttl_min); + handle_holder->mark_protected(); + } else { + handle_holder->mark_delivered(); + } + }); }; server.resource["^/align"]["POST"] = [&](shared_ptr response, @@ -506,11 +611,12 @@ int run_server(Config *config) { if (!config->fnames.size() && anno_graph.wait_for(0s) != std::future_status::ready) throw CurrentlyInitializingError(); // the index is not loaded yet, so we can't process the request - if (!config->fnames.size()) - return process_align_request(content, anno_graph.get()->get_graph(), *config); - - throw std::invalid_argument("Bad request: alignment requests are not yet supported for " - "servers with multiple graphs"); + if (config->fnames.size()) { + throw std::invalid_argument("Bad request: alignment requests are not yet supported for " + "servers with multiple graphs"); + } + return Json::writeString(Json::StreamWriterBuilder(), + process_align_request(content, anno_graph.get()->get_graph(), *config)); }); }; @@ -536,7 +642,7 @@ int run_server(Config *config) { } } } - return root; + return Json::writeString(Json::StreamWriterBuilder(), root); }); }; @@ -595,7 +701,7 @@ int run_server(Config *config) { root["annotation"]["labels"] = get_num_labels(*anno_graph.get()); root["annotation"]["objects"] = static_cast(annotation.num_objects()); } - return root; + return Json::writeString(Json::StreamWriterBuilder(), root); }); }; diff --git a/metagraph/src/cli/server_cache.cpp b/metagraph/src/cli/server_cache.cpp new file mode 100644 index 0000000000..7ba1bdb7c8 --- /dev/null +++ b/metagraph/src/cli/server_cache.cpp @@ -0,0 +1,337 @@ +#include "server_cache.hpp" + +#include +#include +#include +#include + +#include "config/config.hpp" + + +namespace mtg { +namespace cli { + + +// --- Handle ----------------------------------------------------------------- + +ServerQueryCache::Handle::Handle(ServerQueryCache *cache, + std::shared_ptr entry, + std::shared_ptr> producer) + : cache_(cache), + entry_(std::move(entry)), + producer_(std::move(producer)) {} + +ServerQueryCache::Handle::Handle(Handle &&other) noexcept + : cache_(other.cache_), + entry_(std::move(other.entry_)), + producer_(std::move(other.producer_)) { + other.cache_ = nullptr; +} + +ServerQueryCache::Handle & +ServerQueryCache::Handle::operator=(Handle &&other) noexcept { + if (this != &other) { + // Release any current entry first. + if (entry_) + cache_->release_waiter(entry_); + cache_ = other.cache_; + entry_ = std::move(other.entry_); + producer_ = std::move(other.producer_); + other.cache_ = nullptr; + } + return *this; +} + +ServerQueryCache::Handle::~Handle() { + if (!entry_) + return; + // Producer abandoned without publishing (e.g. throw-before-set_result + // along an unexpected path): satisfy the promise with an exception so + // duplicate waiters don't block forever on the shared_future. + if (producer_) { + try { + producer_->set_exception(std::make_exception_ptr( + std::runtime_error("Cache producer abandoned the request"))); + } catch (const std::future_error &) { + // Already satisfied — fine. + } + } + cache_->release_waiter(entry_); +} + +ServerQueryCache::ResultPtr ServerQueryCache::Handle::get() const { + return entry_->future.get(); +} + +void ServerQueryCache::Handle::set_result(std::string response) { + auto ptr = std::make_shared(std::move(response)); + size_t bytes = ptr->size(); + producer_->set_value(ptr); + producer_.reset(); + cache_->on_result_ready(entry_, bytes); +} + +void ServerQueryCache::Handle::set_exception(std::exception_ptr eptr) { + producer_->set_exception(eptr); + producer_.reset(); + // Don't remove the entry — concurrent waiters will receive the + // exception via future.get(). The entry becomes evictable once + // every waiter drops; subsequent identical requests will hit it, + // re-throw the same error, and behave consistently. + cache_->on_result_ready(entry_, /* size = */ 0); +} + +void ServerQueryCache::Handle::mark_delivered() { + cache_->on_delivery(entry_, DeliveryState::DELIVERED); +} + +void ServerQueryCache::Handle::mark_protected() { + cache_->on_delivery(entry_, DeliveryState::PROTECTED); +} + + +// --- ServerQueryCache ------------------------------------------------------- + +ServerQueryCache::ServerQueryCache(size_t max_size_bytes, + std::chrono::nanoseconds protection_ttl) + : max_size_bytes_(max_size_bytes), + protection_ttl_(protection_ttl) {} + +std::pair>, + std::shared_ptr> +ServerQueryCache::make_pending_entry(const std::string &key) { + auto producer = std::make_shared>(); + auto entry = std::make_shared(); + entry->key = key; + entry->future = producer->get_future().share(); + entry->waiters.store(1, std::memory_order_relaxed); + return { std::move(producer), std::move(entry) }; +} + +ServerQueryCache::Handle ServerQueryCache::acquire(const std::string &key) { + if (max_size_bytes_ == 0) { + // Disabled-cache fast path: every caller gets a fresh, detached + // entry. No dedup, no retention, no LRU bookkeeping. + auto [producer, entry] = make_pending_entry(key); + return Handle(this, std::move(entry), std::move(producer)); + } + + std::lock_guard lock(mutex_); + + auto it = map_.find(key); + if (it != map_.end()) { + // Hit. Bump waiters, move to MRU. + auto &entry = it->second; + entry->waiters.fetch_add(1, std::memory_order_relaxed); + touch_lru_locked(entry); + // Sliding retention window: a hit on a PROTECTED entry is + // direct evidence the upstream is still retrying, so we + // refresh ready_at to extend the priority window. + if (entry->delivery.load(std::memory_order_acquire) == DeliveryState::PROTECTED) { + entry->ready_at = std::chrono::steady_clock::now(); + } + return Handle(this, entry, /* producer = */ nullptr); + } + + // Miss. Create a fresh entry, attach it to the cache, and hand the + // promise to the caller (the producer) to fulfil via set_result(). + // Eviction-bookkeeping (size accounting) happens later in + // on_result_ready, when the producer publishes. + auto [producer, entry] = make_pending_entry(key); + entry->in_cache = true; + lru_.push_front(entry); + entry->lru_pos = lru_.begin(); + map_.emplace(key, entry); + return Handle(this, std::move(entry), std::move(producer)); +} + +void ServerQueryCache::release_waiter(const std::shared_ptr &entry) { + int prev = entry->waiters.fetch_sub(1, std::memory_order_acq_rel); + if (prev > 1) + return; // Still has other waiters. + + // Last waiter dropped: this entry — and possibly other entries that + // were skipped during a previous insert because they had waiters — + // are now evictable. Run the sweep so the cache settles back under + // budget without waiting for the next insert. + std::lock_guard lock(mutex_); + evict_under_pressure_locked(); +} + +void ServerQueryCache::on_result_ready(const std::shared_ptr &entry, + size_t size) { + // Called by the producer after publishing the value (set_result) or + // exception (set_exception). Records the moment the result became + // available, registers its size in the byte budget, and runs the + // size-pressure sweep — the just-inserted entry itself is protected + // from eviction in this pass because its producer's Handle is still + // alive (waiters ≥ 1). + std::lock_guard lock(mutex_); + if (!entry->in_cache) + return; // Detached entry from the disabled-cache path. + entry->ready_at = std::chrono::steady_clock::now(); + entry->size_bytes = size; + total_size_bytes_ += size; + evict_under_pressure_locked(); +} + +void ServerQueryCache::on_delivery(const std::shared_ptr &entry, + DeliveryState state) { + // Called from the on_sent async-write callback exactly once per + // delivery attempt. The caller passes either DELIVERED (write + // succeeded) or PROTECTED (write failed → keep the response cached + // with priority for the configured retry window). + if (state == DeliveryState::DELIVERED) { + // Sink: once any delivery has succeeded, the entry is DELIVERED + // forever. A later mark_protected() (e.g. from a duplicate + // request whose delivery dropped) is a no-op below. + entry->delivery.store(DeliveryState::DELIVERED, std::memory_order_release); + return; + } + // state == PROTECTED. CAS-loop the transition so we don't downgrade + // a DELIVERED entry to PROTECTED. + DeliveryState current = entry->delivery.load(std::memory_order_acquire); + bool transitioned_to_protected = false; + while (current != DeliveryState::DELIVERED) { + if (entry->delivery.compare_exchange_weak(current, DeliveryState::PROTECTED, + std::memory_order_acq_rel)) { + transitioned_to_protected = true; + break; + } + } + if (transitioned_to_protected) { + // ready_at is normally already set by on_result_ready (called + // from set_result/set_exception); this is just a defensive + // arm for the producer-abandoned path, where ~Handle satisfies + // the promise without going through on_result_ready. + // Subsequent PROTECTED→PROTECTED transitions are no-ops here — + // the sliding window is refreshed on each cache hit in acquire(). + std::lock_guard lock(mutex_); + if (entry->ready_at.time_since_epoch().count() == 0) + entry->ready_at = std::chrono::steady_clock::now(); + } +} + +void ServerQueryCache::touch_lru_locked(const std::shared_ptr &entry) { + if (!entry->in_cache) + return; + if (entry->lru_pos != lru_.begin()) + lru_.splice(lru_.begin(), lru_, entry->lru_pos); + entry->lru_pos = lru_.begin(); +} + +void ServerQueryCache::evict_under_pressure_locked() { + auto now = std::chrono::steady_clock::now(); + auto is_within_protection_window = [&](const Entry &e) { + // PROTECTED entries within the retention window have *higher* + // priority than DELIVERED ones — they're sacrificed only when + // there's nothing else to give up. + return e.delivery.load(std::memory_order_acquire) == DeliveryState::PROTECTED + && e.ready_at.time_since_epoch().count() != 0 + && now - e.ready_at <= protection_ttl_; + }; + + // Walk LRU back→front (oldest first), evict the first waiterless + // entry that satisfies `pred`. Returns whether anything was evicted. + auto evict_one = [&](auto pred) { + auto rit = std::find_if(lru_.rbegin(), lru_.rend(), + [&](const std::shared_ptr &e) { + return e->waiters.load(std::memory_order_relaxed) == 0 && pred(*e); + }); + if (rit == lru_.rend()) + return false; + auto &entry = *rit; + total_size_bytes_ -= entry->size_bytes; + map_.erase(entry->key); + entry->in_cache = false; + lru_.erase(std::next(rit).base()); + return true; + }; + + // Pass 1: sacrifice DELIVERED and out-of-window PROTECTED entries + // first; in-window PROTECTED entries are held back. + while (total_size_bytes_ > max_size_bytes_ + && evict_one([&](const Entry &e) { return !is_within_protection_window(e); })) { + } + // Pass 2: only triggered when in-window PROTECTED entries are the + // only waiterless candidates left (e.g. cache flooded with retries + // for many distinct failed requests). Sacrifice them in LRU order + // so the cache stays bounded. + while (total_size_bytes_ > max_size_bytes_ + && evict_one([](const Entry &) { return true; })) { + } +} + +size_t ServerQueryCache::size_bytes() const { + std::lock_guard lock(mutex_); + return total_size_bytes_; +} + +size_t ServerQueryCache::entry_count() const { + std::lock_guard lock(mutex_); + return map_.size(); +} + +bool ServerQueryCache::contains(const std::string &key) const { + std::lock_guard lock(mutex_); + return map_.count(key) != 0; +} + + +// --- key construction ------------------------------------------------------- + +std::string make_search_cache_key(const Json::Value &json, + const Config &server_config, + const std::string &graph_identity) { + // Build a stable string key out of every input that affects the + // semantic result. The FASTA body (the bulk of the request) is + // collapsed to a hash; flag-sized values are inlined verbatim. + // + // Per-request overrides are pulled directly from the JSON the same + // way process_search_request resolves them, so flags that select + // different query modes (e.g. abundance_sum vs query_coords) give + // different keys for the same FASTA. + // + // NOTE: keep this resolution logic in sync with the JSON-to-Config + // mapping in `process_search_request` (server.cpp). + const std::string &fasta = json["FASTA"].asString(); + size_t fasta_hash = std::hash{}(fasta); + + double discovery_fraction + = json.get("discovery_fraction", server_config.discovery_fraction).asDouble(); + double min_exact_match + = json.get("min_exact_match", server_config.alignment_min_exact_match).asDouble(); + double max_nodes_per_seq_char + = json.get("max_num_nodes_per_seq_char", + server_config.alignment_max_nodes_per_seq_char).asDouble(); + int top_labels + = json.get("top_labels", server_config.num_top_labels).asInt(); + + int query_mode; + if (json.get("query_coords", false).asBool()) { + query_mode = COORDS; + } else if (json.get("query_counts", false).asBool()) { + query_mode = COUNTS; + } else if (json.get("with_signature", false).asBool()) { + query_mode = SIGNATURE; + } else if (json.get("abundance_sum", false).asBool()) { + query_mode = COUNTS_SUM; + } else { + query_mode = MATCHES; + } + + std::ostringstream oss; + oss << "g:" << graph_identity << '|' + << "m:" << query_mode << '|' + << "df:" << discovery_fraction << '|' + << "ame:" << min_exact_match << '|' + << "amn:" << max_nodes_per_seq_char << '|' + << "tl:" << top_labels << '|' + << "al:" << (json.get("align", false).asBool() ? '1' : '0') << '|' + << "fl:" << fasta.size() << '|' + << "fh:" << fasta_hash; + return oss.str(); +} + +} // namespace cli +} // namespace mtg diff --git a/metagraph/src/cli/server_cache.hpp b/metagraph/src/cli/server_cache.hpp new file mode 100644 index 0000000000..53dbe8eacd --- /dev/null +++ b/metagraph/src/cli/server_cache.hpp @@ -0,0 +1,251 @@ +#ifndef __SERVER_CACHE_HPP__ +#define __SERVER_CACHE_HPP__ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + + +namespace mtg { +namespace cli { + +class Config; + +/** + * A bounded result cache for the `/search` server endpoint. + * + * Goals + * ----- + * 1. Dedup concurrent identical requests so they share one computation. + * 2. Survive connection drops: an upstream retry after a TCP failure + * hits the cached result instead of recomputing a multi-minute query. + * 3. Bound memory while still favoring entries that look like they're + * being retried by a disconnected upstream. + * + * Per-request lifecycle + * --------------------- + * 1. The handler calls `acquire(key)` and gets a `Handle`. waiters++. + * 2. On a miss the handler computes the result and calls + * `Handle::set_result(...)` (or `set_exception(...)` on a thrown + * compute). The `shared_future` underneath unblocks any concurrent + * duplicate caller that arrived during the computation. + * On a hit the handler simply reads the cached value via + * `Handle::get()`. + * 3. The response is written to the wire. The async-write completion + * callback (`on_sent(error_code)`) calls + * ec == 0 → `Handle::mark_delivered()` + * ec != 0 → `Handle::mark_protected()` + * 4. `~Handle()` decrements waiters; the last drop allows the entry + * to be considered for eviction. + * + * Entry states + * ------------ + * PENDING Computation in flight. waiters ≥ 1. Never evicted. + * DELIVERED At least one delivery succeeded. Sink state — a later + * `mark_protected()` is a no-op. Evicted only by LRU under + * size pressure. + * PROTECTED No delivery has succeeded yet. The entry has elevated + * retention priority for `protection_ttl` (default 2 h). + * Every cache hit refreshes the window (sliding TTL — the + * hit itself is direct evidence the upstream is still + * retrying). After 2 h of no hits the entry graduates to + * the main cache: state stays PROTECTED but it loses + * priority and ages out via normal LRU. + * + * Eviction policy (size-pressure sweep, two passes) + * ------------------------------------------------- + * Pass 1: walk LRU back→front, evict the oldest waiterless entry that + * is NOT PROTECTED-within-window. + * Pass 2: still over budget? Walk again, evict any waiterless entry + * (PROTECTED-within-window included). Keeps the cache bounded + * when PROTECTED entries dominate. + * Entries with waiters > 0 are never evicted. + * + * Cache value + * ----------- + * The serialized JSON response body, kept as a `std::string` (held + * via shared_ptr so concurrent waiters/readers can share it). Caching + * the serialized form rather than `Json::Value` (a) makes the byte + * budget reflect the entry's actual memory footprint instead of the + * much larger Json::Value tree, and (b) lets cache hits skip + * serialization entirely. The server's `verbose_output` flag is + * process-constant (set on startup), so the cached body is consistent + * across all hits within a process; only Accept-Encoding-driven + * compression is reapplied per request. + */ +class ServerQueryCache { + public: + enum class DeliveryState { PENDING, DELIVERED, PROTECTED }; + + // The cache value is the serialized JSON response body. + using ResultPtr = std::shared_ptr; + using ResultFuture = std::shared_future; + + private: + // Defined here so `Handle` can refer to `Entry` before its full body. + struct Entry; + + public: + /** + * RAII handle returned by `acquire()`. The handle holds one waiter + * count on the entry and drops it in the destructor. + * + * On a miss (is_miss() == true) this caller is the producer and + * must call exactly one of `set_result(...)` / `set_exception(...)` + * to release any concurrent duplicate waiters. On a hit just call + * `get()` to read the cached value. + * + * After the response has been written to the wire, the on_sent + * callback is expected to call mark_delivered() or mark_protected() + * exactly once. + */ + class Handle { + public: + Handle() = default; + Handle(const Handle &) = delete; + Handle &operator=(const Handle &) = delete; + Handle(Handle &&) noexcept; + Handle &operator=(Handle &&) noexcept; + ~Handle(); + + // True if this handle holds an entry (i.e. acquire() ran). + explicit operator bool() const { return static_cast(entry_); } + + // True if this caller is responsible for computing the result. + bool is_miss() const { return static_cast(producer_); } + + // Wait for the result. Throws whatever the producer set via + // set_exception(). On a miss must be called only after + // set_result/set_exception, otherwise it deadlocks the producer. + ResultPtr get() const; + + // Producer-only: publish the result (or thrown exception) to + // all current and future waiters of this entry's shared_future. + void set_result(std::string response); + void set_exception(std::exception_ptr eptr); + + // Record the eventual delivery outcome of the response. + // + // mark_delivered() is a sink: once any delivery has succeeded + // the entry stays DELIVERED, and a later mark_protected() on + // the same entry (e.g. from a duplicate request whose delivery + // dropped) is a no-op — a successful delivery is a permanent + // fact. + // + // mark_protected() puts the entry under the priority-retention + // window (or extends it if already there). The window timer is + // also refreshed on each cache hit, since a hit is itself + // evidence the upstream is still retrying. + void mark_delivered(); + void mark_protected(); + + private: + friend class ServerQueryCache; + + Handle(ServerQueryCache *cache, + std::shared_ptr entry, + std::shared_ptr> producer); + + ServerQueryCache *cache_ = nullptr; + std::shared_ptr entry_; + std::shared_ptr> producer_; // non-null on miss only + }; + + explicit ServerQueryCache(size_t max_size_bytes, + std::chrono::nanoseconds protection_ttl = std::chrono::hours(2)); + + /** + * Look up `key` and bump its waiter count. + * + * On a miss the returned handle's `is_miss()` is true and the + * caller must produce the result. + * + * On a hit of a PROTECTED entry, the entry's retention window is + * also refreshed — the hit is direct evidence the upstream is + * still retrying. + * + * If the cache is disabled (max_size_bytes == 0) every call is a + * miss and the entry is not retained after the handle drops. + */ + Handle acquire(const std::string &key); + + // Diagnostics / tests. + size_t size_bytes() const; + size_t entry_count() const; + bool contains(const std::string &key) const; + std::chrono::nanoseconds protection_ttl() const { return protection_ttl_; } + + private: + struct Entry { + std::string key; + // Hinge of the dedup mechanism: producer publishes via promise, + // every waiter (including late arrivals) reads via .get(). + ResultFuture future; + // RAII-managed by Handle: producer + each reader contributes 1. + // An entry with waiters > 0 is never evicted. + std::atomic waiters{0}; + // Set when the producer publishes (set_result / set_exception + // → on_result_ready), and refreshed on every cache hit while + // the entry is in PROTECTED state (sliding retention window). + // Used to compute "within protection_ttl" for eviction priority. + std::chrono::steady_clock::time_point ready_at{}; + std::atomic delivery{DeliveryState::PENDING}; + // Byte size of the published response — duplicated here so the + // eviction sweep can sum/subtract it without touching the + // shared_future. Set once in on_result_ready. + size_t size_bytes = 0; + // Iterator into ServerQueryCache::lru_; valid iff in_cache. + std::list>::iterator lru_pos{}; + // false means this entry is detached (evicted, or made by the + // disabled-cache fast path); accounting/eviction skip it. + bool in_cache = false; + }; + + // Allocate a new Entry + paired promise/future. Caller decides + // whether to attach it to the cache (in_cache, lru_, map_) or use + // it as a one-shot detached entry (disabled-cache fast path). + std::pair>, std::shared_ptr> + make_pending_entry(const std::string &key); + + void release_waiter(const std::shared_ptr &entry); + void on_result_ready(const std::shared_ptr &entry, size_t size); + void on_delivery(const std::shared_ptr &entry, DeliveryState state); + void evict_under_pressure_locked(); + void touch_lru_locked(const std::shared_ptr &entry); + + const size_t max_size_bytes_; + const std::chrono::nanoseconds protection_ttl_; + + mutable std::mutex mutex_; + std::unordered_map> map_; + // LRU: front = MRU, back = LRU candidate. + std::list> lru_; + size_t total_size_bytes_ = 0; +}; + + +/** + * Build the canonical cache key for a `/search` request given the parsed + * JSON body, the server's startup Config, and the resolved graph identity. + * + * Per-request overrides (discovery_fraction, query_mode flags, top_labels, + * etc.) are read from `json` directly, mirroring how `process_search_request` + * resolves them. The key includes only inputs that affect the *semantic* + * result; it deliberately excludes formatting flags like `verbose_output` + * and `Accept-Encoding: deflate`. + */ +std::string make_search_cache_key(const Json::Value &json, + const Config &server_config, + const std::string &graph_identity); + +} // namespace cli +} // namespace mtg + +#endif // __SERVER_CACHE_HPP__ diff --git a/metagraph/src/cli/server_utils.cpp b/metagraph/src/cli/server_utils.cpp index acedb591c2..62f7e9cd55 100644 --- a/metagraph/src/cli/server_utils.cpp +++ b/metagraph/src/cli/server_utils.cpp @@ -84,20 +84,21 @@ std::string json_str_with_error_msg(const std::string &msg) { void process_request(std::shared_ptr &response, const std::shared_ptr &request, size_t request_id, - const std::function &process) { + const std::function &process, + std::function on_sent) { logger->info("[Server] {} request {} from {}", request->path, request_id, request->remote_endpoint().address().to_string()); Timer timer; - // Retrieve string: std::string content = request->content.string(); SimpleWeb::CaseInsensitiveMultimap header({ { "Content-Type", "application/json" } }); SimpleWeb::StatusCode status; std::string ret; try { - // Return JSON string status = SimpleWeb::StatusCode::success_ok; - ret = Json::writeString(Json::StreamWriterBuilder(), process(content)); + // process() is expected to return the response body already + // serialized as JSON (a std::string). + ret = process(content); if (is_compression_requested(request)) { ret = compress_string(ret); header.insert(std::make_pair("Content-Encoding", "deflate")); @@ -119,10 +120,12 @@ void process_request(std::shared_ptr &response, } double processing_time = timer.elapsed(); response->write(status, ret, header); + response->send(std::move(on_sent)); logger->info("[Server] Request {} processing time: {:.3f} sec, response size: {:.1f} KB, " "finished in {:.3f} sec", request_id, processing_time, (double)ret.size() / 1000, timer.elapsed()); } + } // namespace cli } // namespace mtg diff --git a/metagraph/src/cli/server_utils.hpp b/metagraph/src/cli/server_utils.hpp index 5cfce14bf5..33e9a658e1 100644 --- a/metagraph/src/cli/server_utils.hpp +++ b/metagraph/src/cli/server_utils.hpp @@ -9,10 +9,28 @@ namespace cli { using HttpServer = SimpleWeb::Server; +/** + * Run the request handler `process` and write its response to the client. + * + * `process` returns the response body as an already-serialized JSON + * string. Handlers that build `Json::Value` should serialize it inline + * (e.g. via `Json::writeString(StreamWriterBuilder(), value)`); /search + * uses this form so cache hits can return the serialized body directly + * without re-serializing the JSON tree on every hit. + * + * If `on_sent` is non-null it is invoked with the asio error_code + * returned by the async network write: + * ec == 0 → bytes left this host without a socket-level error + * ec != 0 → the connection failed before/during delivery + * Used by /search to drive the result cache: a non-zero ec marks the + * cached entry PROTECTED so an upstream's retry hits cache instead of + * recomputing. + */ void process_request(std::shared_ptr &response, const std::shared_ptr &request, size_t request_id, - const std::function &process); + const std::function &process, + std::function on_sent = nullptr); class CurrentlyInitializingError : public std::runtime_error { public: diff --git a/metagraph/tests/cli/test_server_cache.cpp b/metagraph/tests/cli/test_server_cache.cpp new file mode 100644 index 0000000000..1adb061816 --- /dev/null +++ b/metagraph/tests/cli/test_server_cache.cpp @@ -0,0 +1,350 @@ +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "cli/server_cache.hpp" + + +namespace { + +using namespace mtg::cli; +using namespace std::chrono_literals; + +// Build a placeholder response string of the requested byte size. The +// cache value is the serialized JSON body — for the eviction-policy +// tests we only care about its size, so the contents are arbitrary. +std::string response_of_size(size_t bytes) { + return std::string(bytes, '\0'); +} + + +TEST(ServerQueryCache, MissThenHit) { + ServerQueryCache cache(1ull << 20); + auto h1 = cache.acquire("k1"); + ASSERT_TRUE(h1.is_miss()); + h1.set_result(response_of_size(100)); + auto val1 = h1.get(); + ASSERT_NE(val1, nullptr); + + auto h2 = cache.acquire("k1"); + EXPECT_FALSE(h2.is_miss()); + auto val2 = h2.get(); + EXPECT_EQ(val1.get(), val2.get()); // both handles share the same response object +} + + +TEST(ServerQueryCache, ConcurrentDeduplication) { + ServerQueryCache cache(1ull << 20); + + // Producer thread acquires first, holds the handle, then publishes. + std::atomic compute_count = 0; + std::promise waiter_can_start; + auto producer = std::thread([&]() { + auto h = cache.acquire("dedup"); + ASSERT_TRUE(h.is_miss()); + compute_count.fetch_add(1); + // Let the waiter run while we're still computing. + waiter_can_start.set_value(); + std::this_thread::sleep_for(50ms); + h.set_result(response_of_size(42)); + h.get(); // ensure the future is observed before destruction + }); + + waiter_can_start.get_future().wait(); + auto h = cache.acquire("dedup"); + // Second arrival must NOT be a miss — first computation owns the slot. + EXPECT_FALSE(h.is_miss()); + auto val = h.get(); // blocks until producer publishes + EXPECT_EQ(val->size(), 42u); + producer.join(); + + // Single computation despite two acquires. + EXPECT_EQ(compute_count.load(), 1); +} + + +TEST(ServerQueryCache, EvictionUnderSizePressure) { + ServerQueryCache cache(/* max */ 100); + + { + auto h1 = cache.acquire("a"); + h1.set_result(response_of_size(60)); + } + { + auto h2 = cache.acquire("b"); + h2.set_result(response_of_size(60)); + } + // Total 120 > 100 → LRU "a" should be evicted. + EXPECT_FALSE(cache.contains("a")); + EXPECT_TRUE(cache.contains("b")); + EXPECT_LE(cache.size_bytes(), 100u); +} + + +TEST(ServerQueryCache, NeverEvictsEntryWithWaiters) { + ServerQueryCache cache(/* max */ 100); + + // "a" has a live waiter — should never be evicted even under pressure. + auto h_a = cache.acquire("a"); + h_a.set_result(response_of_size(60)); + + // Insert "b" and let it become evictable (waiters drop to 0). + { + auto h_b = cache.acquire("b"); + h_b.set_result(response_of_size(60)); + } + // Total at this point is 120 > 100, but "b" had waiters when its + // result was published, so it survived the insert-time eviction pass. + // Adding "c" triggers another eviction sweep: "a" is protected, + // "b" is now waiterless, so "b" is the LRU victim. + { + auto h_c = cache.acquire("c"); + h_c.set_result(response_of_size(60)); + } + EXPECT_TRUE(cache.contains("a")); + EXPECT_FALSE(cache.contains("b")); +} + + +TEST(ServerQueryCache, ProtectedWithinWindowPreservedOverDelivered) { + // Pass 1 should evict the DELIVERED entry before touching the + // PROTECTED entry that's still inside its retention window. + ServerQueryCache cache(/* max */ 100, /* protection_ttl */ 1h); + + { + auto h_failed = cache.acquire("failed"); + h_failed.set_result(response_of_size(60)); + h_failed.mark_protected(); + } + { + auto h_ok = cache.acquire("ok"); + h_ok.set_result(response_of_size(40)); + h_ok.mark_delivered(); + } + // Both waiterless, total=100 (== max). Force a touch under pressure. + { + auto h_extra = cache.acquire("extra"); + h_extra.set_result(response_of_size(40)); // 100 + 40 = 140 > 100 + h_extra.mark_delivered(); + } + EXPECT_TRUE(cache.contains("failed")); // protected + EXPECT_FALSE(cache.contains("ok")); // sacrificed first + EXPECT_TRUE(cache.contains("extra")); +} + + +TEST(ServerQueryCache, ProtectedEntriesEvictedUnderHeavyPressure) { + // When all waiterless entries are PROTECTED-within-window and the cache is + // overfull, pass 2 must still evict — otherwise the cache grows without + // bound. + ServerQueryCache cache(/* max */ 100, /* protection_ttl */ 1h); + + auto store_failed = [&](const std::string &key, size_t bytes) { + auto h = cache.acquire(key); + h.set_result(response_of_size(bytes)); + h.mark_protected(); + }; + + store_failed("f1", 60); + store_failed("f2", 60); // 60 + 60 = 120 > 100 + // f1 was protected during f2's set_result (f2 had waiters); f1 is now + // the only waiterless entry, and pass 2 lets us evict it despite TTL. + EXPECT_FALSE(cache.contains("f1")); + EXPECT_TRUE(cache.contains("f2")); + EXPECT_LE(cache.size_bytes(), 100u); +} + + +TEST(ServerQueryCache, ProtectedPastWindowGraduatesToMainCache) { + // After `protection_ttl` of no retry hits, a PROTECTED entry stays + // in the cache but loses its priority — it's evictable like a + // DELIVERED entry, only under size pressure (no proactive removal). + ServerQueryCache cache(/* max */ 100, /* protection_ttl */ 50ms); + + { + auto h = cache.acquire("aged"); + h.set_result(response_of_size(60)); + h.mark_protected(); + } + std::this_thread::sleep_for(120ms); + // Past the window — entry stays in cache. + EXPECT_TRUE(cache.contains("aged")); + + // Insert a DELIVERED entry that triggers size pressure. The aged + // PROTECTED entry is no longer prioritized → evictable in pass 1 + // at normal LRU priority. + { + auto h = cache.acquire("new"); + h.set_result(response_of_size(60)); // 60 + 60 = 120 > 100 + h.mark_delivered(); + } + EXPECT_FALSE(cache.contains("aged")); + EXPECT_TRUE(cache.contains("new")); +} + + +TEST(ServerQueryCache, ProtectedHitRefreshesPriorityWindow) { + // A cache hit on a PROTECTED entry refreshes its retention TTL — + // the window is sliding, tracking the upstream's retry pattern. + ServerQueryCache cache(/* max */ 100, /* protection_ttl */ 100ms); + + { + auto h = cache.acquire("retried"); + h.set_result(response_of_size(60)); + h.mark_protected(); + } + // Sacrificial DELIVERED entry: filling cache to capacity so a later + // insert triggers pressure with a waiterless eviction candidate + // available (otherwise pass 1 would have nothing to evict and pass 2 + // would falsely sacrifice the PROTECTED entry we're trying to test). + { + auto h = cache.acquire("filler"); + h.set_result(response_of_size(40)); // total = 100, at cap + h.mark_delivered(); + } + + std::this_thread::sleep_for(60ms); + // Mid-window retry: cache hit refreshes ready_at on "retried". + { + auto h = cache.acquire("retried"); + ASSERT_FALSE(h.is_miss()); + h.get(); + } + std::this_thread::sleep_for(60ms); + // Elapsed = 120ms (past the original 100ms TTL). With the sliding + // window refresh we're only 60ms into the new window — still protected. + + // Insert an intruder that triggers pressure. Pass 1 should sacrifice + // "filler" (DELIVERED, waiterless) and leave "retried" alone. + { + auto h = cache.acquire("intruder"); + h.set_result(response_of_size(40)); // total during insert: 140 > 100 + h.mark_delivered(); + } + EXPECT_TRUE(cache.contains("retried")); // protected by refreshed window + EXPECT_FALSE(cache.contains("filler")); // sacrificed in pass 1 + EXPECT_TRUE(cache.contains("intruder")); +} + + +TEST(ServerQueryCache, ProtectedWithinWindowIsKept) { + ServerQueryCache cache(/* max */ 1ull << 20, /* protection_ttl */ 1h); + + { + auto h = cache.acquire("kept"); + h.set_result(response_of_size(10)); + h.mark_protected(); + } + EXPECT_TRUE(cache.contains("kept")); + auto h = cache.acquire("kept"); + EXPECT_FALSE(h.is_miss()); + EXPECT_EQ(h.get()->size(), 10u); +} + + +TEST(ServerQueryCache, DeliveredIsSinkStateAgainstLaterFailure) { + ServerQueryCache cache(/* max */ 1ull << 20, /* protection_ttl */ 50ms); + + // First delivery succeeds — entry becomes DELIVERED. + { + auto h = cache.acquire("k"); + h.set_result(response_of_size(10)); + h.mark_delivered(); + } + // Second acquirer of the same entry has a delivery failure. + { + auto h = cache.acquire("k"); + ASSERT_FALSE(h.is_miss()); + h.mark_protected(); // must NOT downgrade DELIVERED → PROTECTED + } + // Wait past the (very short) protection TTL. With the sink semantics + // the entry stays DELIVERED so the eviction sweep treats it as a + // normal LRU citizen rather than a downgraded PROTECTED-past-window + // entry — but we never get to test that distinction, the assertion + // below is just "still in cache". + std::this_thread::sleep_for(120ms); + { + // Trigger an eviction sweep via release_waiter. + auto h = cache.acquire("touch"); + h.set_result(response_of_size(10)); + } + EXPECT_TRUE(cache.contains("k")); +} + + +TEST(ServerQueryCache, DeliveredEntryStaysUntilSizePressure) { + ServerQueryCache cache(/* max */ 1ull << 20, /* protection_ttl */ 1ms); + + { + auto h = cache.acquire("ok"); + h.set_result(response_of_size(10)); + h.mark_delivered(); + } + std::this_thread::sleep_for(20ms); + // Even though the protection TTL has elapsed, DELIVERED entries are + // immune — only LRU pressure can evict them. + auto h = cache.acquire("ok"); + EXPECT_FALSE(h.is_miss()); +} + + +TEST(ServerQueryCache, DisabledCacheRecomputesEveryTime) { + ServerQueryCache cache(/* max */ 0); + + auto h1 = cache.acquire("x"); + EXPECT_TRUE(h1.is_miss()); + h1.set_result(response_of_size(10)); + + // Disabled cache: nothing is retained, second acquire is also a miss. + auto h2 = cache.acquire("x"); + EXPECT_TRUE(h2.is_miss()); + EXPECT_EQ(cache.entry_count(), 0u); + EXPECT_FALSE(cache.contains("x")); +} + + +TEST(ServerQueryCache, ProducerExceptionPropagatesToWaiters) { + ServerQueryCache cache(1ull << 20); + + std::promise waiter_ready; + auto producer = std::thread([&]() { + auto h = cache.acquire("boom"); + ASSERT_TRUE(h.is_miss()); + waiter_ready.set_value(); + std::this_thread::sleep_for(50ms); + h.set_exception(std::make_exception_ptr(std::runtime_error("kaboom"))); + }); + + waiter_ready.get_future().wait(); + auto h = cache.acquire("boom"); + EXPECT_FALSE(h.is_miss()); + EXPECT_THROW(h.get(), std::runtime_error); + producer.join(); +} + + +TEST(ServerQueryCache, AbandonedProducerUnblocksWaiters) { + ServerQueryCache cache(1ull << 20); + + std::promise waiter_ready; + auto producer = std::thread([&]() { + auto h = cache.acquire("orphan"); + ASSERT_TRUE(h.is_miss()); + waiter_ready.set_value(); + std::this_thread::sleep_for(20ms); + // Drop the handle without ever calling set_result — destructor must + // satisfy the promise with an exception so waiters don't hang. + }); + + waiter_ready.get_future().wait(); + auto h = cache.acquire("orphan"); + EXPECT_FALSE(h.is_miss()); + EXPECT_THROW(h.get(), std::runtime_error); + producer.join(); +} + + +} // namespace