c6dfae5661
Seastar is an external library from the point of view of ScyllaDB, so should be included with angle brackets. Closes scylladb/scylladb#27947
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
|