scylladb/vector_search/vector_store_client.cc

/*
 * Copyright (C) 2025-present ScyllaDB
 */

/*
 * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
 */

#include "vector_store_client.hh"
#include "dns.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 = vector_search::vector_store_client::ann_error;
using configuration_exception = exceptions::configuration_exception;
using duration = lowres_clock::duration;
using vs_vector = vector_search::vector_store_client::vs_vector;
using limit = vector_search::vector_store_client::limit;
using host_name = vector_search::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 = vector_search::vector_store_client::port_number;
using primary_key = vector_search::vector_store_client::primary_key;
using primary_keys = vector_search::vector_store_client::primary_keys;
using service_reply_format_error = vector_search::vector_store_client::service_reply_format_error;
using tcp_keepalive_params = net::tcp_keepalive_params;
using time_point = lowres_clock::time_point;

/// 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 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 vector_search {

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;
    gate client_producer_gate;
    condition_variable refresh_client_cv;
    milliseconds wait_for_client_timeout = WAIT_FOR_CLIENT_TIMEOUT;
    unsigned http_request_retries = HTTP_REQUEST_RETRIES;
    sequential_producer<lw_shared_ptr<http_client>> client_producer;
    dns dns;

    impl(utils::config_file::named_value<sstring> cfg)
        : uri_observer(cfg.observe([this](seastar::sstring uri) {
            try {
                handle_uri_changed(get_host_port(uri));
            } catch (const configuration_exception& e) {
                vslogger.error("Failed to parse Vector Store service URI: {}", e.what());
                handle_uri_changed(std::nullopt);
            }
        }))
        , _host_port(get_host_port(cfg()))
        , client_producer([&]() -> future<lw_shared_ptr<http_client>> {
            return try_with_gate(client_producer_gate, [this] -> future<lw_shared_ptr<http_client>> {
                dns.trigger_refresh();
                co_await wait_for_signal(refresh_client_cv, lowres_clock::now() + wait_for_client_timeout);
                co_return current_client;
            });
        })
        , dns(vslogger, _host_port ? std::optional<seastar::sstring>{_host_port->host} : std::nullopt, [this](auto const& addr) -> future<> {
            co_await handle_address_changed(addr);
        }) {
    }

    void handle_uri_changed(std::optional<host_port> uri) {
        _host_port = std::move(uri);
        if (current_client) {
            old_clients.emplace_back(current_client);
            current_client = nullptr;
        }
        dns.host(_host_port ? std::optional<seastar::sstring>{_host_port->host} : std::nullopt);
    }

    auto handle_address_changed(const std::optional<inet_address>& addr) -> future<> {
        if (current_client) {
            old_clients.emplace_back(current_client);
            current_client = nullptr;
        }
        if (addr && _host_port) {
            current_client = make_lw_shared<http_client>(*_host_port, std::move(*addr));
        }
        refresh_client_cv.broadcast();
        co_await cleanup_old_clients();
    }

    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;
    }

    /// 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;
            }

            dns.trigger_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() {
    _impl->dns.start_background_tasks();
}

auto vector_store_client::stop() -> future<> {
    _impl->refresh_client_cv.signal();
    co_await _impl->client_producer_gate.close();
    co_await _impl->dns.stop();
    co_await _impl->cleanup_old_clients();
    co_await _impl->cleanup_current_client();
}

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->dns.trigger_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 vector_search