mirrors/scylladb
1
0
Fork 0
mirror of https://github.com/scylladb/scylladb.git synced 2026-04-24 18:40:38 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
b18b052d2615bfafb64d3b09b0333c0d15fbe3fa
scylladb/service/vector_store_client.cc
Karol Nowacki eedf506be5 vector_store_client: Rename vector_store_uri to vector_store_primary_uri 
The configuration setting vector_store_uri is renamed to
vector_store_primary_uri according to the final design.
In the future, the vector_store_secondary_uri setting will
be introduced.

This setting now also accepts a comma-separated list of URIs to prepare
for future support for redundancy and load balancing. Currently, only the
first URI in the list is used.

This change must be included before the next release.
Otherwise, users will be affected by a breaking change.

References: VECTOR-187

Closes scylladb/scylladb#26033
2025-09-21 16:33:10 +03:00

654 lines
24 KiB
C++
Raw Blame History

/*
* Copyright (C) 2025-present ScyllaDB
*/
/*
* SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
*/
#include "vector_store_client.hh"
#include "cql3/statements/select_statement.hh"
#include "cql3/type_json.hh"
#include "db/config.hh"
#include "exceptions/exceptions.hh"
#include "utils/sequential_producer.hh"
#include "dht/i_partitioner.hh"
#include "keys/keys.hh"
#include "utils/rjson.hh"
#include "schema/schema.hh"
#include <charconv>
#include <exception>
#include <fmt/ranges.h>
#include <regex>
#include <seastar/core/sstring.hh>
#include <seastar/coroutine/as_future.hh>
#include <seastar/coroutine/exception.hh>
#include <seastar/http/client.hh>
#include <seastar/http/request.hh>
#include <seastar/net/dns.hh>
#include <seastar/net/inet_address.hh>
#include <seastar/net/socket_defs.hh>
#include <seastar/util/short_streams.hh>
namespace {
using namespace std::chrono_literals;
using ann_error = service::vector_store_client::ann_error;
using configuration_exception = exceptions::configuration_exception;
using duration = lowres_clock::duration;
using vs_vector = service::vector_store_client::vs_vector;
using limit = service::vector_store_client::limit;
using host_name = service::vector_store_client::host_name;
using http_path = sstring;
using inet_address = seastar::net::inet_address;
using json_content = sstring;
using milliseconds = std::chrono::milliseconds;
using operation_type = httpd::operation_type;
using port_number = service::vector_store_client::port_number;
using primary_key = service::vector_store_client::primary_key;
using primary_keys = service::vector_store_client::primary_keys;
using service_reply_format_error = service::vector_store_client::service_reply_format_error;
using tcp_keepalive_params = net::tcp_keepalive_params;
using time_point = lowres_clock::time_point;
// Wait time before retrying after an exception occurred
constexpr auto EXCEPTION_OCCURED_WAIT = std::chrono::seconds(5);
// Minimum interval between dns name refreshes
constexpr auto DNS_REFRESH_INTERVAL = std::chrono::seconds(5);
/// Timeout for waiting for a new client to be available
constexpr auto WAIT_FOR_CLIENT_TIMEOUT = std::chrono::seconds(5);
/// How many retries to do for HTTP requests
constexpr auto HTTP_REQUEST_RETRIES = 3;
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
logging::logger vslogger("vector_store_client");
auto parse_port(std::string const& port_txt) -> std::optional<port_number> {
auto port = port_number{};
auto [ptr, ec] = std::from_chars(&*port_txt.begin(), &*port_txt.end(), port);
if (*ptr != '\0' || ec != std::errc{}) {
return std::nullopt;
}
return port;
}
struct host_port {
host_name host;
port_number port;
};
auto parse_service_uri(std::string_view uri) -> std::optional<host_port> {
constexpr auto URI_REGEX = R"(^http:\/\/([a-z0-9._-]+):([0-9]+)$)";
auto const uri_regex = std::regex(URI_REGEX);
auto uri_match = std::smatch{};
auto uri_txt = std::string(uri);
if (!std::regex_match(uri_txt, uri_match, uri_regex) || uri_match.size() != 3) {
return {};
}
auto host = uri_match[1].str();
auto port = parse_port(uri_match[2].str());
if (!port) {
return {};
}
return {{host, *port}};
}
/// Wait for a timeout ar abort signal.
auto wait_for_timeout(duration timeout, abort_source& as) -> future<bool> {
auto result = co_await coroutine::as_future(sleep_abortable(timeout, as));
if (result.failed()) {
auto err = result.get_exception();
if (as.abort_requested()) {
co_return false;
}
co_await coroutine::return_exception_ptr(std::move(err));
}
co_return true;
}
/// Wait for a condition variable to be signaled or timeout.
auto wait_for_signal(condition_variable& cv, time_point timeout) -> future<void> {
auto result = co_await coroutine::as_future(cv.wait(timeout));
if (result.failed()) {
auto err = result.get_exception();
if (try_catch<condition_variable_timed_out>(err) != nullptr) {
co_return;
}
co_await coroutine::return_exception_ptr(std::move(err));
}
co_return;
}
auto get_key_column_value(const rjson::value& item, std::size_t idx, const column_definition& column) -> std::expected<bytes, ann_error> {
auto const& column_name = column.name_as_text();
auto const* keys_obj = rjson::find(item, column_name);
if (keys_obj == nullptr) {
vslogger.error("Vector Store returned invalid JSON: missing key column '{}'", column_name);
return std::unexpected{service_reply_format_error{}};
}
if (!keys_obj->IsArray()) {
vslogger.error("Vector Store returned invalid JSON: key column '{}' is not an array", column_name);
return std::unexpected{service_reply_format_error{}};
}
auto const& keys_arr = keys_obj->GetArray();
if (keys_arr.Size() <= idx) {
vslogger.error("Vector Store returned invalid JSON: key column '{}' array too small", column_name);
return std::unexpected{service_reply_format_error{}};
}
auto const& key = keys_arr[idx];
return from_json_object(*column.type, key);
}
auto pk_from_json(rjson::value const& item, std::size_t idx, schema_ptr const& schema) -> std::expected<partition_key, ann_error> {
std::vector<bytes> raw_pk;
for (const column_definition& cdef : schema->partition_key_columns()) {
auto raw_value = get_key_column_value(item, idx, cdef);
if (!raw_value) {
return std::unexpected{raw_value.error()};
}
raw_pk.emplace_back(*raw_value);
}
return partition_key::from_exploded(raw_pk);
}
auto ck_from_json(rjson::value const& item, std::size_t idx, schema_ptr const& schema) -> std::expected<clustering_key_prefix, ann_error> {
if (schema->clustering_key_size() == 0) {
return clustering_key_prefix::make_empty();
}
std::vector<bytes> raw_ck;
for (const column_definition& cdef : schema->clustering_key_columns()) {
auto raw_value = get_key_column_value(item, idx, cdef);
if (!raw_value) {
return std::unexpected{raw_value.error()};
}
raw_ck.emplace_back(*raw_value);
}
return clustering_key_prefix::from_exploded(raw_ck);
}
auto write_ann_json(vs_vector vs_vector, limit limit) -> json_content {
return seastar::format(R"({{"vector":[{}],"limit":{}}})", fmt::join(vs_vector, ","), limit);
}
auto read_ann_json(rjson::value const& json, schema_ptr const& schema) -> std::expected<primary_keys, ann_error> {
if (!json.HasMember("primary_keys")) {
vslogger.error("Vector Store returned invalid JSON: missing 'primary_keys'");
return std::unexpected{service_reply_format_error{}};
}
auto const& keys_json = json["primary_keys"];
if (!keys_json.IsObject()) {
vslogger.error("Vector Store returned invalid JSON: 'primary_keys' is not an object");
return std::unexpected{service_reply_format_error{}};
}
if (!json.HasMember("distances")) {
vslogger.error("Vector Store returned invalid JSON: missing 'distances'");
return std::unexpected{service_reply_format_error{}};
}
auto const& distances_json = json["distances"];
if (!distances_json.IsArray()) {
vslogger.error("Vector Store returned invalid JSON: 'distances' is not an array");
return std::unexpected{service_reply_format_error{}};
}
auto const& distances_arr = json["distances"].GetArray();
auto size = distances_arr.Size();
auto keys = primary_keys{};
for (auto idx = 0U; idx < size; ++idx) {
auto pk = pk_from_json(keys_json, idx, schema);
if (!pk) {
return std::unexpected{pk.error()};
}
auto ck = ck_from_json(keys_json, idx, schema);
if (!ck) {
return std::unexpected{ck.error()};
}
keys.push_back(primary_key{dht::decorate_key(*schema, *pk), *ck});
}
return std::move(keys);
}
class client_connection_factory : public http::experimental::connection_factory {
socket_address _addr;
public:
explicit client_connection_factory(socket_address addr)
: _addr(addr) {
}
future<connected_socket> make([[maybe_unused]] abort_source* as) override {
auto socket = co_await seastar::connect(_addr, {}, transport::TCP);
socket.set_nodelay(true);
socket.set_keepalive_parameters(tcp_keepalive_params{
.idle = 60s,
.interval = 60s,
.count = 10,
});
socket.set_keepalive(true);
co_return socket;
}
};
class http_client {
host_port _host_port;
inet_address _addr;
http::experimental::client impl;
public:
http_client(host_port host_port_, inet_address addr)
: _host_port(std::move(host_port_))
, _addr(std::move(addr))
, impl(std::make_unique<client_connection_factory>(socket_address(addr, _host_port.port))) {
}
bool connects_to(inet_address const& a, port_number p) const {
return _addr == a && _host_port.port == p;
}
seastar::future<> make_request(
operation_type method, http_path path, std::optional<json_content> content, http::experimental::client::reply_handler&& handle, abort_source* as) {
auto req = http::request::make(method, _host_port.host, std::move(path));
if (content) {
req.write_body("json", std::move(*content));
}
return impl.make_request(std::move(req), std::move(handle), std::nullopt, as);
}
seastar::future<> close() {
return impl.close();
}
const inet_address& addr() const {
return _addr;
}
};
bool should_vector_store_service_be_disabled(std::vector<sstring> const& uris) {
return uris.empty() || uris[0].empty();
}
auto get_host_port(std::string_view uris_csv) -> std::optional<host_port> {
auto uris = utils::split_comma_separated_list(uris_csv);
if (should_vector_store_service_be_disabled(uris)) {
vslogger.info("Vector Store service URIs are empty, disabling Vector Store service");
return std::nullopt;
}
auto uri = uris[0];
auto parsed = parse_service_uri(uri);
if (!parsed) {
throw configuration_exception(fmt::format("Invalid Vector Store service URI: {}", uri));
}
vslogger.info("Vector Store service URI is set to '{}'", uri);
return *parsed;
}
sstring response_content_to_sstring(const std::vector<temporary_buffer<char>>& buffers) {
sstring result;
for (const auto& buf : buffers) {
result.append(buf.get(), buf.size());
}
return result;
}
} // namespace
namespace service {
struct vector_store_client::impl {
utils::observer<sstring> uri_observer;
lw_shared_ptr<http_client> current_client;
std::vector<lw_shared_ptr<http_client>> old_clients;
std::optional<host_port> _host_port;
time_point last_dns_refresh;
gate tasks_gate;
gate client_producer_gate;
condition_variable refresh_cv;
condition_variable refresh_client_cv;
abort_source abort_refresh;
milliseconds dns_refresh_interval = DNS_REFRESH_INTERVAL;
milliseconds wait_for_client_timeout = WAIT_FOR_CLIENT_TIMEOUT;
unsigned http_request_retries = HTTP_REQUEST_RETRIES;
std::function<future<std::optional<inet_address>>(sstring const&)> dns_resolver;
sequential_producer<lw_shared_ptr<http_client>> client_producer;
impl(utils::config_file::named_value<sstring> cfg)
: uri_observer(cfg.observe([this](std::string_view uri) {
try {
_host_port = get_host_port(uri);
trigger_dns_refresh();
} catch (const configuration_exception& e) {
vslogger.error("Failed to parse Vector Store service URI: {}", e.what());
_host_port = std::nullopt;
}
}))
, _host_port(get_host_port(cfg()))
, dns_resolver([](auto const& host) -> future<std::optional<inet_address>> {
auto addr = co_await coroutine::as_future(net::dns::resolve_name(host));
if (addr.failed()) {
auto err = addr.get_exception();
if (try_catch<std::system_error>(err) != nullptr) {
co_return std::nullopt;
}
co_await coroutine::return_exception_ptr(std::move(err));
}
co_return co_await std::move(addr);
})
, client_producer([this]() -> future<lw_shared_ptr<http_client>> {
return try_with_gate(client_producer_gate, [this] -> future<lw_shared_ptr<http_client>> {
trigger_dns_refresh();
co_await wait_for_signal(refresh_client_cv, lowres_clock::now() + wait_for_client_timeout);
co_return current_client;
});
}) {
}
auto is_disabled() const -> bool {
return !bool{_host_port};
}
auto host() const -> std::expected<host_name, disabled> {
if (is_disabled()) {
return std::unexpected{disabled{}};
}
return _host_port->host;
}
auto port() const -> std::expected<port_number, disabled> {
if (is_disabled()) {
return std::unexpected{disabled{}};
}
return _host_port->port;
}
/// Refresh the http client with a new address resolved from the DNS name.
/// If the DNS resolution fails, the current client is set to nullptr.
/// If the address is the same as the current one, do nothing.
/// Old clients are saved for later cleanup in a specific task.
auto refresh_addr() -> future<> {
if (is_disabled()) {
current_client = nullptr;
co_return;
}
auto [host, port] = *_host_port;
auto new_addr = co_await dns_resolver(host);
if (!new_addr) {
current_client = nullptr;
co_return;
}
// Check if the new address and port is the same as the current one
if (current_client && current_client->connects_to(*new_addr, port)) {
co_return;
}
old_clients.emplace_back(current_client);
current_client = make_lw_shared<http_client>(*_host_port, std::move(*new_addr));
}
/// A task for refreshing the vector store http client.
auto refresh_addr_task() -> future<> {
for (;;) {
auto exception_occured = false;
try {
if (abort_refresh.abort_requested()) {
break;
}
// Do not refresh the service address too often
auto now = lowres_clock::now();
auto current_duration = now - last_dns_refresh;
if (current_duration > dns_refresh_interval) {
last_dns_refresh = now;
co_await refresh_addr();
} else {
// Wait till the end of the refreshing interval
if (co_await wait_for_timeout(dns_refresh_interval - current_duration, abort_refresh)) {
continue;
}
// If the wait was aborted, we stop refreshing
break;
}
if (abort_refresh.abort_requested()) {
break;
}
// new client is available
refresh_client_cv.broadcast();
co_await cleanup_old_clients();
co_await refresh_cv.when();
} catch (const std::exception& e) {
vslogger.error("Vector Store Client refresh task failed: {}", e.what());
exception_occured = true;
} catch (...) {
vslogger.error("Vector Store Client refresh task failed with unknown exception");
exception_occured = true;
}
if (exception_occured) {
// If an exception occurred, we wait for the next signal to refresh the address
co_await wait_for_timeout(EXCEPTION_OCCURED_WAIT, abort_refresh);
}
}
co_await cleanup_old_clients();
co_await cleanup_current_client();
}
/// Request a DNS refresh in the specific task.
void trigger_dns_refresh() {
refresh_cv.signal();
}
/// Cleanup current client
auto cleanup_current_client() -> future<> {
if (current_client) {
co_await current_client->close();
}
current_client = nullptr;
}
/// Cleanup old clients that are no longer used.
auto cleanup_old_clients() -> future<> {
// iterate over old clients and close them. There is a co_await in the loop
// so we need to use [] accessor and copying clients to avoid dangling references of iterators.
// NOLINTNEXTLINE(modernize-loop-convert)
for (auto it = 0U; it < old_clients.size(); ++it) {
auto& client = old_clients[it];
if (client && client.owned()) {
auto client_cloned = client;
co_await client_cloned->close();
client_cloned = nullptr;
}
}
std::erase_if(old_clients, [](auto const& client) {
return !client;
});
}
using get_client_error = std::variant<aborted, addr_unavailable, disabled>;
/// Get the current http client or wait for a new one to be available.
auto get_client(abort_source& as) -> future<std::expected<lw_shared_ptr<http_client>, get_client_error>> {
if (is_disabled()) {
co_return std::unexpected{disabled{}};
}
if (current_client) {
co_return current_client;
}
auto current_client = co_await coroutine::as_future(client_producer(as));
if (current_client.failed()) {
auto err = current_client.get_exception();
if (as.abort_requested()) {
co_return std::unexpected{aborted{}};
}
co_await coroutine::return_exception_ptr(std::move(err));
}
auto client = co_await std::move(current_client);
if (!client) {
co_return std::unexpected{addr_unavailable{}};
}
co_return client;
}
struct make_request_response {
http::reply::status_type status; ///< The HTTP status of the response.
std::vector<temporary_buffer<char>> content; ///< The content of the response.
};
using make_request_error = std::variant<aborted, addr_unavailable, service_unavailable, disabled>;
auto make_request(operation_type method, http_path path, std::optional<json_content> content, abort_source& as)
-> future<std::expected<make_request_response, make_request_error>> {
auto resp = make_request_response{.status = http::reply::status_type::ok, .content = std::vector<temporary_buffer<char>>()};
for (auto retries = 0; retries < HTTP_REQUEST_RETRIES; ++retries) {
auto client = co_await get_client(as);
if (!client) {
co_return std::unexpected{std::visit(
[](auto&& err) {
return make_request_error{err};
},
client.error())};
}
auto result = co_await coroutine::as_future(client.value()->make_request(
method, std::move(path), std::move(content),
[&resp](http::reply const& reply, input_stream<char> body) -> future<> {
resp.status = reply._status;
resp.content = co_await util::read_entire_stream(body);
},
&as));
if (result.failed()) {
auto err = result.get_exception();
if (as.abort_requested()) {
co_return std::unexpected{aborted{}};
}
if (try_catch<std::system_error>(err) == nullptr) {
co_await coroutine::return_exception_ptr(std::move(err));
}
// std::system_error means that the server is unavailable, so we retry
} else {
co_return resp;
}
trigger_dns_refresh();
}
co_return std::unexpected{service_unavailable{}};
}
};
vector_store_client::vector_store_client(config const& cfg)
: _impl(std::make_unique<impl>(cfg.vector_store_primary_uri)) {
}
vector_store_client::~vector_store_client() = default;
void vector_store_client::start_background_tasks() {
/// start the background task to refresh the service address
(void)try_with_gate(_impl->tasks_gate, [this] {
return _impl->refresh_addr_task();
}).handle_exception([](std::exception_ptr eptr) {
on_internal_error_noexcept(vslogger, format("The Vector Store Client refresh task failed: {}", eptr));
});
}
auto vector_store_client::stop() -> future<> {
_impl->abort_refresh.request_abort();
_impl->refresh_cv.signal();
_impl->refresh_client_cv.signal();
co_await _impl->tasks_gate.close();
co_await _impl->client_producer_gate.close();
}
auto vector_store_client::is_disabled() const -> bool {
return _impl->is_disabled();
}
auto vector_store_client::host() const -> std::expected<host_name, disabled> {
return _impl->host();
}
auto vector_store_client::port() const -> std::expected<port_number, disabled> {
return _impl->port();
}
auto vector_store_client::ann(keyspace_name keyspace, index_name name, schema_ptr schema, vs_vector vs_vector, limit limit, abort_source& as)
-> future<std::expected<primary_keys, ann_error>> {
if (is_disabled()) {
vslogger.error("Disabled Vector Store while calling ann");
co_return std::unexpected{disabled{}};
}
auto path = format("/api/v1/indexes/{}/{}/ann", keyspace, name);
auto content = write_ann_json(std::move(vs_vector), limit);
auto resp = co_await _impl->make_request(operation_type::POST, std::move(path), std::move(content), as);
if (!resp) {
co_return std::unexpected{std::visit(
[](auto&& err) {
return ann_error{err};
},
resp.error())};
}
if (resp->status != status_type::ok) {
vslogger.error("Vector Store returned error: HTTP status {}: {}", resp->status,
seastar::value_of([&resp] {return response_content_to_sstring(resp->content);}));
co_return std::unexpected{service_error{resp->status}};
}
try {
co_return read_ann_json(rjson::parse(std::move(resp->content)), schema);
} catch (const rjson::error& e) {
vslogger.error("Vector Store returned invalid JSON: {}", e.what());
co_return std::unexpected{service_reply_format_error{}};
}
}
void vector_store_client_tester::set_dns_refresh_interval(vector_store_client& vsc, std::chrono::milliseconds interval) {
vsc._impl->dns_refresh_interval = interval;
}
void vector_store_client_tester::set_wait_for_client_timeout(vector_store_client& vsc, std::chrono::milliseconds timeout) {
vsc._impl->wait_for_client_timeout = timeout;
}
void vector_store_client_tester::set_http_request_retries(vector_store_client& vsc, unsigned retries) {
vsc._impl->http_request_retries = retries;
}
void vector_store_client_tester::set_dns_resolver(vector_store_client& vsc, std::function<future<std::optional<inet_address>>(sstring const&)> resolver) {
vsc._impl->dns_resolver = std::move(resolver);
}
void vector_store_client_tester::trigger_dns_resolver(vector_store_client& vsc) {
vsc._impl->trigger_dns_refresh();
}
auto vector_store_client_tester::resolve_hostname(vector_store_client& vsc, abort_source& as) -> future<std::optional<inet_address>> {
auto client = co_await vsc._impl->get_client(as);
if (!client) {
co_return std::nullopt;
}
co_return client.value()->addr();
}
} // namespace service
Reference in New Issue View Git Blame Copy Permalink