Merge pull request #797 from Altinity/feature/antalya-25.3/rendezvous_hashing

25.3 Antalya port of #709, #760 - Rendezvous hashing

Author: Enmk

src/Storages/IStorageCluster.cpp

@@ -117,7 +117,7 @@ void ReadFromCluster::createExtension(const ActionsDAG::Node * predicate, size_t
     if (extension)
         return;
 
-    extension = storage->getTaskIteratorExtension(predicate, context, number_of_replicas);
+    extension = storage->getTaskIteratorExtension(predicate, context, cluster);
 }
 
 /// The code executes on initiator
@@ -196,7 +196,7 @@ void ReadFromCluster::initializePipeline(QueryPipelineBuilder & pipeline, const
     if (current_settings[Setting::max_parallel_replicas] > 1)
         max_replicas_to_use = std::min(max_replicas_to_use, current_settings[Setting::max_parallel_replicas].value);
 
-    createExtension(nullptr, max_replicas_to_use);
+    createExtension(nullptr);
 
     for (const auto & shard_info : cluster->getShardsInfo())
     {

src/Storages/IStorageCluster.h

@@ -35,7 +35,10 @@ class IStorageCluster : public IStorage
 
     ClusterPtr getCluster(ContextPtr context) const;
     /// Query is needed for pruning by virtual columns (_file, _path)
-    virtual RemoteQueryExecutor::Extension getTaskIteratorExtension(const ActionsDAG::Node * predicate, const ContextPtr & context, size_t number_of_replicas) const = 0;
+    virtual RemoteQueryExecutor::Extension getTaskIteratorExtension(
+        const ActionsDAG::Node * predicate,
+        const ContextPtr & context,
+        ClusterPtr cluster) const = 0;
 
     QueryProcessingStage::Enum getQueryProcessingStage(ContextPtr, QueryProcessingStage::Enum, const StorageSnapshotPtr &, SelectQueryInfo &) const override;
 

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -16,7 +16,6 @@
 #include <Storages/extractTableFunctionFromSelectQuery.h>
 #include <Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.h>
 
-
 namespace DB
 {
 namespace Setting
@@ -177,13 +176,28 @@ void StorageObjectStorageCluster::updateQueryToSendIfNeeded(
 }
 
 RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExtension(
-    const ActionsDAG::Node * predicate, const ContextPtr & local_context, const size_t number_of_replicas) const
+    const ActionsDAG::Node * predicate,
+    const ContextPtr & local_context,
+    ClusterPtr cluster) const
 {
     auto iterator = StorageObjectStorageSource::createFileIterator(
         configuration, configuration->getQuerySettings(local_context), object_storage, /* distributed_processing */false,
         local_context, predicate, {}, virtual_columns, nullptr, local_context->getFileProgressCallback(), /*ignore_archive_globs=*/true, /*skip_object_metadata=*/true);
 
-    auto task_distributor = std::make_shared<StorageObjectStorageStableTaskDistributor>(iterator, number_of_replicas);
+    std::vector<std::string> ids_of_hosts;
+    for (const auto & shard : cluster->getShardsInfo())
+    {
+        if (shard.per_replica_pools.empty())
+            throw Exception(ErrorCodes::LOGICAL_ERROR, "Cluster {} with empty shard {}", cluster->getName(), shard.shard_num);
+        for (const auto & replica : shard.per_replica_pools)
+        {
+            if (!replica)
+                throw Exception(ErrorCodes::LOGICAL_ERROR, "Cluster {}, shard {} with empty node", cluster->getName(), shard.shard_num);
+            ids_of_hosts.push_back(replica->getAddress());
+        }
+    }
+
+    auto task_distributor = std::make_shared<StorageObjectStorageStableTaskDistributor>(iterator, ids_of_hosts);
 
     auto callback = std::make_shared<TaskIterator>(
         [task_distributor](size_t number_of_current_replica) mutable -> String {

src/Storages/ObjectStorage/StorageObjectStorageCluster.h

@@ -24,7 +24,9 @@ class StorageObjectStorageCluster : public IStorageCluster
     std::string getName() const override;
 
     RemoteQueryExecutor::Extension getTaskIteratorExtension(
-        const ActionsDAG::Node * predicate, const ContextPtr & context, size_t number_of_replicas) const override;
+        const ActionsDAG::Node * predicate,
+        const ContextPtr & context,
+        ClusterPtr cluster) const override;
 
     String getPathSample(StorageInMemoryMetadata metadata, ContextPtr context);
 

src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.cpp

@@ -8,9 +8,10 @@ namespace DB
 
 StorageObjectStorageStableTaskDistributor::StorageObjectStorageStableTaskDistributor(
     std::shared_ptr<IObjectIterator> iterator_,
-    size_t number_of_replicas_)
+    std::vector<std::string> ids_of_nodes_)
     : iterator(std::move(iterator_))
-    , connection_to_files(number_of_replicas_)
+    , connection_to_files(ids_of_nodes_.size())
+    , ids_of_nodes(ids_of_nodes_)
     , iterator_exhausted(false)
 {
 }
@@ -37,13 +38,39 @@ std::optional<String> StorageObjectStorageStableTaskDistributor::getNextTask(siz
 
 size_t StorageObjectStorageStableTaskDistributor::getReplicaForFile(const String & file_path)
 {
-    return ConsistentHashing(sipHash64(file_path), connection_to_files.size());
+    size_t nodes_count = ids_of_nodes.size();
+
+    /// Trivial case
+    if (nodes_count < 2)
+        return 0;
+
+    /// Rendezvous hashing
+    size_t best_id = 0;
+    UInt64 best_weight = sipHash64(ids_of_nodes[0] + file_path);
+    for (size_t id = 1; id < nodes_count; ++id)
+    {
+        UInt64 weight = sipHash64(ids_of_nodes[id] + file_path);
+        if (weight > best_weight)
+        {
+            best_weight = weight;
+            best_id = id;
+        }
+    }
+    return best_id;
 }
 
 std::optional<String> StorageObjectStorageStableTaskDistributor::getPreQueuedFile(size_t number_of_current_replica)
 {
     std::lock_guard lock(mutex);
 
+    if (connection_to_files.size() <= number_of_current_replica)
+        throw Exception(
+            ErrorCodes::LOGICAL_ERROR,
+            "Replica number {} is out of range. Expected range: [0, {})",
+            number_of_current_replica,
+            connection_to_files.size()
+        );
+
     auto & files = connection_to_files[number_of_current_replica];
 
     while (!files.empty())

src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.h

@@ -4,6 +4,7 @@
 #include <Common/Logger.h>
 #include <Interpreters/Cluster.h>
 #include <Storages/ObjectStorage/StorageObjectStorageSource.h>
+#include <Storages/ObjectStorageQueue/ObjectStorageQueueSource.h>
 #include <unordered_set>
 #include <vector>
 #include <mutex>
@@ -17,7 +18,7 @@ class StorageObjectStorageStableTaskDistributor
 public:
     StorageObjectStorageStableTaskDistributor(
         std::shared_ptr<IObjectIterator> iterator_,
-        size_t number_of_replicas_);
+        std::vector<std::string> ids_of_nodes_);
 
     std::optional<String> getNextTask(size_t number_of_current_replica);
 
@@ -32,6 +33,8 @@ class StorageObjectStorageStableTaskDistributor
     std::vector<std::vector<String>> connection_to_files;
     std::unordered_set<String> unprocessed_files;
 
+    std::vector<std::string> ids_of_nodes;
+
     std::mutex mutex;
     bool iterator_exhausted = false;
 

src/Storages/StorageDistributed.cpp

@@ -1300,8 +1300,7 @@ std::optional<QueryPipeline> StorageDistributed::distributedWriteFromClusterStor
     const auto cluster = getCluster();
 
     /// Select query is needed for pruining on virtual columns
-    auto number_of_replicas = static_cast<UInt64>(cluster->getShardsInfo().size());
-    auto extension = src_storage_cluster.getTaskIteratorExtension(predicate, local_context, number_of_replicas);
+    auto extension = src_storage_cluster.getTaskIteratorExtension(predicate, local_context, cluster);
 
     /// Here we take addresses from destination cluster and assume source table exists on these nodes
     size_t replica_index = 0;

src/Storages/StorageFileCluster.cpp

@@ -77,7 +77,10 @@ void StorageFileCluster::updateQueryToSendIfNeeded(DB::ASTPtr & query, const Sto
     );
 }
 
-RemoteQueryExecutor::Extension StorageFileCluster::getTaskIteratorExtension(const ActionsDAG::Node * predicate, const ContextPtr & context, const size_t) const
+RemoteQueryExecutor::Extension StorageFileCluster::getTaskIteratorExtension(
+    const ActionsDAG::Node * predicate,
+    const ContextPtr & context,
+    ClusterPtr) const
 {
     auto iterator = std::make_shared<StorageFileSource::FilesIterator>(paths, std::nullopt, predicate, getVirtualsList(), context);
     auto callback = std::make_shared<TaskIterator>([iter = std::move(iterator)](size_t) mutable -> String { return iter->next(); });

src/Storages/StorageFileCluster.h

@@ -27,7 +27,10 @@ class StorageFileCluster : public IStorageCluster
         const ConstraintsDescription & constraints_);
 
     std::string getName() const override { return "FileCluster"; }
-    RemoteQueryExecutor::Extension getTaskIteratorExtension(const ActionsDAG::Node * predicate, const ContextPtr & context, size_t number_of_replicas) const override;
+    RemoteQueryExecutor::Extension getTaskIteratorExtension(
+        const ActionsDAG::Node * predicate,
+        const ContextPtr & context,
+        ClusterPtr) const override;
 
 private:
     void updateQueryToSendIfNeeded(ASTPtr & query, const StorageSnapshotPtr & storage_snapshot, const ContextPtr & context) override;

src/Storages/StorageReplicatedMergeTree.cpp

@@ -6089,8 +6089,7 @@ std::optional<QueryPipeline> StorageReplicatedMergeTree::distributedWriteFromClu
     ContextMutablePtr query_context = Context::createCopy(local_context);
     query_context->increaseDistributedDepth();
 
-    auto number_of_replicas = static_cast<UInt64>(src_cluster->getShardsAddresses().size());
-    auto extension = src_storage_cluster->getTaskIteratorExtension(nullptr, local_context, number_of_replicas);
+    auto extension = src_storage_cluster->getTaskIteratorExtension(nullptr, local_context, src_cluster);
 
     size_t replica_index = 0;
     for (const auto & replicas : src_cluster->getShardsAddresses())