a0e62ef8de
This change introduces a load balancing mechanism for the vector store client. The client can now distribute requests across multiple vector store nodes. The distribution mechanism performs random selection of nodes for each request.
51 lines
1.1 KiB
C++
51 lines
1.1 KiB
C++
/*
|
|
* Copyright (C) 2025-present ScyllaDB
|
|
*/
|
|
|
|
/*
|
|
* SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include "seastar/core/shared_ptr.hh"
|
|
#include <vector>
|
|
|
|
namespace vector_search {
|
|
|
|
template <class T, typename RandomNumberEngine>
|
|
class load_balancer {
|
|
public:
|
|
load_balancer(std::vector<seastar::lw_shared_ptr<T>> container, RandomNumberEngine& g)
|
|
: _container(std::move(container))
|
|
, _g(g) {
|
|
}
|
|
|
|
seastar::lw_shared_ptr<T> next() {
|
|
if (_container.empty()) {
|
|
return nullptr;
|
|
}
|
|
return pop(randomize_index());
|
|
}
|
|
|
|
private:
|
|
using distribution = std::uniform_int_distribution<std::size_t>;
|
|
|
|
size_t randomize_index() {
|
|
return _dist(_g, distribution::param_type(0, _container.size() - 1));
|
|
}
|
|
|
|
seastar::lw_shared_ptr<T> pop(size_t index) {
|
|
auto ret = _container[index];
|
|
std::swap(_container[index], _container.back());
|
|
_container.pop_back();
|
|
return ret;
|
|
}
|
|
|
|
std::vector<seastar::lw_shared_ptr<T>> _container;
|
|
RandomNumberEngine& _g;
|
|
distribution _dist;
|
|
};
|
|
|
|
} // namespace vector_search
|