couchbase 4.2.11-rc.1 → 4.3.0

package/deps/couchbase-cxx-client/core/range_scan_load_balancer.hxx

@@ -0,0 +1,64 @@
+/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
+/*
+ *   Copyright 2024. Couchbase, Inc.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ */
+
+#include "core/logger/logger.hxx"
+#include "core/topology/configuration.hxx"
+
+#include <mutex>
+#include <queue>
+
+namespace couchbase::core
+{
+class range_scan_node_state
+{
+  public:
+    range_scan_node_state(std::queue<std::uint16_t> vbuckets);
+
+    auto fetch_vbucket_id() -> std::optional<std::uint16_t>;
+    void notify_stream_ended();
+    void enqueue_vbucket(std::uint16_t vbucket_id);
+    auto active_stream_count() -> std::uint16_t;
+    auto pending_vbucket_count() -> std::size_t;
+
+  private:
+    std::uint16_t active_stream_count_{ 0 };
+    std::queue<std::uint16_t> pending_vbuckets_{};
+    std::mutex mutex_{};
+};
+
+class range_scan_load_balancer
+{
+  public:
+    range_scan_load_balancer(const topology::configuration::vbucket_map& vbucket_map, std::optional<std::uint64_t> seed = {});
+
+    void seed(std::uint64_t seed);
+
+    /**
+     * Returns the ID of a vbucket that corresponds to the node with the lowest number of active streams.
+     * Returns "std::nullopt" if there are no pending vbuckets
+     */
+    auto select_vbucket() -> std::optional<std::uint16_t>;
+
+    void notify_stream_ended(std::int16_t node_id);
+    void enqueue_vbucket(std::int16_t node_id, std::uint16_t vbucket_id);
+
+  private:
+    std::map<std::int16_t, range_scan_node_state> nodes_{};
+    std::mutex select_vbucket_mutex_{};
+    std::optional<std::uint64_t> seed_{};
+};
+} // namespace couchbase::core