Преглед изворни кода

feat(embeddings): per-origin keep-alive client pool (configurable size) to remove per-call TLS handshakes

Fszontagh пре 1 месец
родитељ
комит
a2686d850d
9 измењених фајлова са 183 додато и 13 уклоњено
  1. 1 1
      api/llms.txt
  2. 2 1
      api/openapi.json
  3. 93 7
      src/embeddings.cpp
  4. 3 2
      src/embeddings.hpp
  5. 2 0
      src/handlers/settings.cpp
  6. 2 1
      src/handlers/vectors.cpp
  7. 3 1
      src/settings.cpp
  8. 1 0
      src/settings.hpp
  9. 76 0
      tests/test_embeddings.cpp

+ 1 - 1
api/llms.txt

@@ -65,7 +65,7 @@ Requires a `kind=vector` collection. Vectors are stored with cosine-similarity i
 All settings endpoints require an admin key.
 
 - `GET  /api/v1/settings` — Return current settings as JSON. `openai_api_key` is masked: the field is absent and `openai_api_key_set` (bool) indicates whether a key is configured.
-- `PUT  /api/v1/settings` `{openai_api_base?, openai_api_key?, default_embedding_model?, cors_origins?, session_ttl_minutes?, webui_enabled?, default_project?, embedding_connect_timeout_sec?, embedding_read_timeout_sec?, embedding_cache_size?, embedding_cache_ttl_sec?, embedding_cache_max_bytes?, embedding_cache_normalize?}` — Merge the supplied fields into the current settings and persist. Omit `openai_api_key` (or pass an empty string) to keep the existing secret. Changes are hot-reloaded immediately; cache settings are applied instantly via `reconfigure()`.
+- `PUT  /api/v1/settings` `{openai_api_base?, openai_api_key?, default_embedding_model?, cors_origins?, session_ttl_minutes?, webui_enabled?, default_project?, embedding_connect_timeout_sec?, embedding_read_timeout_sec?, embedding_cache_size?, embedding_cache_ttl_sec?, embedding_cache_max_bytes?, embedding_cache_normalize?, embedding_client_pool_size?}` — Merge the supplied fields into the current settings and persist. Omit `openai_api_key` (or pass an empty string) to keep the existing secret. Changes are hot-reloaded immediately; cache settings are applied instantly via `reconfigure()`.
 
 ## Stats
 

+ 2 - 1
api/openapi.json

@@ -852,7 +852,8 @@
                   "embedding_cache_size": { "type": "integer", "minimum": 0, "description": "Maximum number of cached embedding entries (default 4096)" },
                   "embedding_cache_ttl_sec": { "type": "integer", "minimum": 0, "description": "Cache entry TTL in seconds; 0 = no expiry (default 86400)" },
                   "embedding_cache_max_bytes": { "type": "integer", "minimum": 0, "description": "Maximum total bytes of cached vectors (default 268435456 = 256 MB)" },
-                  "embedding_cache_normalize": { "type": "boolean", "description": "When true, text is trimmed and lower-cased before cache lookup (default false)" }
+                  "embedding_cache_normalize": { "type": "boolean", "description": "When true, text is trimmed and lower-cased before cache lookup (default false)" },
+                  "embedding_client_pool_size": { "type": "integer", "minimum": 1, "description": "Number of keep-alive HTTP clients pooled per upstream origin (default 4)" }
                 }
               }
             }

+ 93 - 7
src/embeddings.cpp

@@ -1,7 +1,93 @@
 #include "embeddings.hpp"
 #include "errors.hpp"
 #include <httplib.h>
+#include <condition_variable>
+#include <map>
+#include <memory>
+#include <mutex>
+#include <vector>
 namespace svapi {
+namespace {
+
+// A pool of N independent keep-alive httplib::Clients for a single upstream origin.
+// Each Client owns one TCP/TLS connection; a thread must have EXCLUSIVE use of a
+// Client between checkout and checkin (httplib::Client is not safe for concurrent
+// use by multiple threads). Checkout blocks until a free Client is available.
+class ClientPool {
+public:
+    ClientPool(const std::string& origin, const std::string& apiKey,
+               uint32_t connectTimeoutSec, uint32_t readTimeoutSec, uint32_t poolSize) {
+        uint32_t n = poolSize < 1 ? 1 : poolSize;
+        clients_.reserve(n);
+        free_.reserve(n);
+        for (uint32_t i = 0; i < n; ++i) {
+            auto cli = std::make_unique<httplib::Client>(origin);
+            cli->set_keep_alive(true);
+            cli->set_connection_timeout(connectTimeoutSec);
+            cli->set_read_timeout(readTimeoutSec);
+            cli->enable_server_certificate_verification(true);
+            if (!apiKey.empty()) cli->set_bearer_token_auth(apiKey);
+            free_.push_back(cli.get());
+            clients_.push_back(std::move(cli));
+        }
+    }
+
+    // RAII lease: borrows one Client on construction (blocking) and returns it on
+    // destruction, so the Client is freed even if the POST throws.
+    class Lease {
+    public:
+        explicit Lease(ClientPool& pool) : pool_(pool) {
+            std::unique_lock<std::mutex> lk(pool_.mu_);
+            pool_.cv_.wait(lk, [&] { return !pool_.free_.empty(); });
+            cli_ = pool_.free_.back();
+            pool_.free_.pop_back();
+        }
+        ~Lease() {
+            {
+                std::lock_guard<std::mutex> lk(pool_.mu_);
+                pool_.free_.push_back(cli_);
+            }
+            pool_.cv_.notify_one();
+        }
+        Lease(const Lease&) = delete;
+        Lease& operator=(const Lease&) = delete;
+        httplib::Client& client() { return *cli_; }
+    private:
+        ClientPool& pool_;
+        httplib::Client* cli_;
+    };
+
+private:
+    std::mutex mu_;
+    std::condition_variable cv_;
+    std::vector<std::unique_ptr<httplib::Client>> clients_;  // owns the Clients
+    std::vector<httplib::Client*> free_;                     // currently-idle Clients
+};
+
+// Process-wide registry keyed on everything that affects a connection's behavior.
+// A settings hot-reload that changes any of these fields naturally routes to a
+// freshly-built pool; the stale pool simply ages out of use (no eviction — YAGNI).
+ClientPool& poolFor(const std::string& origin, const std::string& apiKey,
+                    uint32_t connectTimeoutSec, uint32_t readTimeoutSec, uint32_t poolSize) {
+    static std::mutex regMu;
+    static std::map<std::string, std::unique_ptr<ClientPool>> registry;
+    // '\x1f' (unit separator) cannot appear in origins/keys/decimals → unambiguous key.
+    std::string key = origin + '\x1f' + apiKey + '\x1f' +
+                      std::to_string(connectTimeoutSec) + '\x1f' +
+                      std::to_string(readTimeoutSec) + '\x1f' +
+                      std::to_string(poolSize);
+    std::lock_guard<std::mutex> lk(regMu);
+    auto it = registry.find(key);
+    if (it == registry.end()) {
+        it = registry.emplace(key, std::make_unique<ClientPool>(
+                                       origin, apiKey, connectTimeoutSec, readTimeoutSec, poolSize))
+                 .first;
+    }
+    return *it->second;
+}
+
+} // namespace
+
 std::vector<float> parseEmbeddingResponse(const nlohmann::json& body) {
     if (!body.contains("data") || !body["data"].is_array() || body["data"].empty())
         throw ApiError(ErrCode::Unprocessable, "embedding_failed", "embedding response missing data[]");
@@ -23,17 +109,17 @@ EmbeddingEndpoint splitEmbeddingBase(const std::string& apiBase) {
 }
 
 EmbeddingClient::EmbeddingClient(std::string apiBase, std::string apiKey,
-                                 uint32_t connectTimeoutSec, uint32_t readTimeoutSec)
+                                 uint32_t connectTimeoutSec, uint32_t readTimeoutSec,
+                                 uint32_t poolSize)
     : apiBase_(std::move(apiBase)), apiKey_(std::move(apiKey)),
-      connectTimeoutSec_(connectTimeoutSec), readTimeoutSec_(readTimeoutSec) {}
+      connectTimeoutSec_(connectTimeoutSec), readTimeoutSec_(readTimeoutSec),
+      poolSize_(poolSize) {}
 std::vector<float> EmbeddingClient::embed(const std::string& model, const std::string& text) {
     EmbeddingEndpoint ep = splitEmbeddingBase(apiBase_);
-    httplib::Client cli(ep.origin);
-    cli.set_connection_timeout(connectTimeoutSec_); cli.set_read_timeout(readTimeoutSec_);
-    cli.enable_server_certificate_verification(true);
-    if (!apiKey_.empty()) cli.set_bearer_token_auth(apiKey_);
+    ClientPool& pool = poolFor(ep.origin, apiKey_, connectTimeoutSec_, readTimeoutSec_, poolSize_);
+    ClientPool::Lease lease(pool);  // EXCLUSIVE use of one keep-alive Client until scope exit
     nlohmann::json req{{"model", model}, {"input", text}};
-    auto res = cli.Post(ep.path + "/v1/embeddings", req.dump(), "application/json");
+    auto res = lease.client().Post(ep.path + "/v1/embeddings", req.dump(), "application/json");
     if (!res) throw ApiError(ErrCode::Unavailable, "embedding_unreachable", "could not reach embedding endpoint");
     if (res->status != 200)
         throw ApiError(ErrCode::Unprocessable, "embedding_failed",

+ 3 - 2
src/embeddings.hpp

@@ -15,10 +15,11 @@ EmbeddingEndpoint splitEmbeddingBase(const std::string& apiBase);
 class EmbeddingClient {
 public:
     EmbeddingClient(std::string apiBase, std::string apiKey,
-                    uint32_t connectTimeoutSec = 10, uint32_t readTimeoutSec = 60);
+                    uint32_t connectTimeoutSec = 10, uint32_t readTimeoutSec = 60,
+                    uint32_t poolSize = 4);
     std::vector<float> embed(const std::string& model, const std::string& text);
 private:
     std::string apiBase_, apiKey_;
-    uint32_t connectTimeoutSec_, readTimeoutSec_;
+    uint32_t connectTimeoutSec_, readTimeoutSec_, poolSize_;
 };
 }

+ 2 - 0
src/handlers/settings.cpp

@@ -49,6 +49,8 @@ void registerSettingsRoutes(ApiServer& s) {
             updated.embeddingCacheMaxBytes = body["embedding_cache_max_bytes"].get<uint64_t>();
         if (body.contains("embedding_cache_normalize"))
             updated.embeddingCacheNormalize = body["embedding_cache_normalize"].get<bool>();
+        if (body.contains("embedding_client_pool_size"))
+            updated.embeddingClientPoolSize = body["embedding_client_pool_size"].get<uint32_t>();
         d->settings.save(updated);
         sendJson(res, 200, {{"ok", true}});
     });

+ 2 - 1
src/handlers/vectors.cpp

@@ -32,7 +32,8 @@ std::vector<float> resolveVector(ServerDeps* d, const CollectionMeta& meta, cons
             v = std::move(*hit);
         } else {
             EmbeddingClient emb(snap->openaiApiBase, snap->openaiApiKey,
-                                snap->embeddingConnectTimeoutSec, snap->embeddingReadTimeoutSec);
+                                snap->embeddingConnectTimeoutSec, snap->embeddingReadTimeoutSec,
+                                snap->embeddingClientPoolSize);
             v = emb.embed(model, text);
             embeddingCache().put(key, v);  // copy, not move: v is still needed for the dimension check below
         }

+ 3 - 1
src/settings.cpp

@@ -18,6 +18,7 @@ Settings Settings::fromJson(const nlohmann::json& j) {
     s.embeddingCacheTtlSec       = j.value("embedding_cache_ttl_sec",     s.embeddingCacheTtlSec);
     s.embeddingCacheMaxBytes     = j.value("embedding_cache_max_bytes",   s.embeddingCacheMaxBytes);
     s.embeddingCacheNormalize    = j.value("embedding_cache_normalize",   s.embeddingCacheNormalize);
+    s.embeddingClientPoolSize    = j.value("embedding_client_pool_size",  s.embeddingClientPoolSize);
     return s;
 }
 
@@ -31,7 +32,8 @@ nlohmann::json Settings::toJson() const {
             {"embedding_cache_size",        embeddingCacheSize},
             {"embedding_cache_ttl_sec",     embeddingCacheTtlSec},
             {"embedding_cache_max_bytes",   embeddingCacheMaxBytes},
-            {"embedding_cache_normalize",   embeddingCacheNormalize}};
+            {"embedding_cache_normalize",   embeddingCacheNormalize},
+            {"embedding_client_pool_size",  embeddingClientPoolSize}};
 }
 
 } // namespace svapi

+ 1 - 0
src/settings.hpp

@@ -20,6 +20,7 @@ struct Settings {
     uint32_t    embeddingCacheTtlSec       = 86400;      // 0 = no expiry
     uint64_t    embeddingCacheMaxBytes     = 268435456;  // 256 MB
     bool        embeddingCacheNormalize    = false;      // opt-in trim+lowercase
+    uint32_t    embeddingClientPoolSize    = 4;          // keep-alive clients per upstream origin
 
     static Settings fromJson(const nlohmann::json& j);
     nlohmann::json  toJson() const;

+ 76 - 0
tests/test_embeddings.cpp

@@ -3,7 +3,10 @@
 #include "embedding_cache.hpp"
 #include "errors.hpp"
 #include <httplib.h>
+#include <atomic>
+#include <chrono>
 #include <thread>
+#include <vector>
 using namespace svapi;
 
 TEST(Embeddings, ParsesWellFormedResponse) {
@@ -53,6 +56,79 @@ TEST(Embeddings, ClientHonorsBasePathPrefix) {
     mock.stop(); t.join();
 }
 
+TEST(Embeddings, PoolRunsRequestsConcurrently) {
+    // A serialized impl (single Client + mutex) would pin maxInFlight at 1.
+    // A real pool of N clients lets up to N requests overlap on independent sockets.
+    std::atomic<int> inFlight{0};
+    std::atomic<int> maxInFlight{0};
+    httplib::Server mock;
+    mock.Post("/v1/embeddings", [&](const httplib::Request&, httplib::Response& res) {
+        int now = ++inFlight;
+        int prev = maxInFlight.load();
+        while (now > prev && !maxInFlight.compare_exchange_weak(prev, now)) {}
+        std::this_thread::sleep_for(std::chrono::milliseconds(180));
+        --inFlight;
+        res.set_content(R"({"data":[{"embedding":[1,2,3,4]}]})", "application/json");
+    });
+    int port = mock.bind_to_any_port("127.0.0.1");
+    mock.set_keep_alive_max_count(1000);
+    std::thread srv([&]{ mock.listen_after_bind(); });
+    while (!mock.is_running()) std::this_thread::sleep_for(std::chrono::milliseconds(5));
+
+    const std::string origin = "http://127.0.0.1:" + std::to_string(port);
+    constexpr int K = 4;
+    std::vector<std::thread> workers;
+    std::atomic<int> ok{0};
+    for (int i = 0; i < K; ++i) {
+        workers.emplace_back([&]{
+            EmbeddingClient cli(origin, "", 10, 60, /*poolSize=*/4);
+            auto v = cli.embed("m", "t");
+            if (v.size() == 4u && v[3] == 4.0f) ++ok;
+        });
+    }
+    for (auto& w : workers) w.join();
+    mock.stop(); srv.join();
+
+    EXPECT_EQ(ok.load(), K);                 // all requests succeeded
+    EXPECT_GE(maxInFlight.load(), 2);        // proves overlap — pool did NOT serialize
+}
+
+TEST(Embeddings, PoolSizeOneSerializes) {
+    // poolSize=1 means a single keep-alive client → requests cannot overlap.
+    std::atomic<int> inFlight{0};
+    std::atomic<int> maxInFlight{0};
+    httplib::Server mock;
+    mock.Post("/v1/embeddings", [&](const httplib::Request&, httplib::Response& res) {
+        int now = ++inFlight;
+        int prev = maxInFlight.load();
+        while (now > prev && !maxInFlight.compare_exchange_weak(prev, now)) {}
+        std::this_thread::sleep_for(std::chrono::milliseconds(120));
+        --inFlight;
+        res.set_content(R"({"data":[{"embedding":[1,2,3,4]}]})", "application/json");
+    });
+    int port = mock.bind_to_any_port("127.0.0.1");
+    mock.set_keep_alive_max_count(1000);
+    std::thread srv([&]{ mock.listen_after_bind(); });
+    while (!mock.is_running()) std::this_thread::sleep_for(std::chrono::milliseconds(5));
+
+    const std::string origin = "http://127.0.0.1:" + std::to_string(port);
+    constexpr int K = 4;
+    std::vector<std::thread> workers;
+    std::atomic<int> ok{0};
+    for (int i = 0; i < K; ++i) {
+        workers.emplace_back([&]{
+            EmbeddingClient cli(origin, "", 10, 60, /*poolSize=*/1);
+            auto v = cli.embed("m", "t");
+            if (v.size() == 4u) ++ok;
+        });
+    }
+    for (auto& w : workers) w.join();
+    mock.stop(); srv.join();
+
+    EXPECT_EQ(ok.load(), K);
+    EXPECT_EQ(maxInFlight.load(), 1);        // serialized by design (single client)
+}
+
 // ---- EmbeddingCache tests -----------------------------------------------
 
 TEST(EmbeddingCache, ReturnsCachedVectorOnExactKey) {