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

src/QueryPipeline/RemoteQueryExecutor.cpp

@@ -738,8 +738,12 @@ void RemoteQueryExecutor::processReadTaskRequest()
     if (!extension || !extension->task_iterator)
         throw Exception(ErrorCodes::LOGICAL_ERROR, "Distributed task iterator is not initialized");
 
+    if (!extension->replica_info)
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Replica info is not initialized");
+
     ProfileEvents::increment(ProfileEvents::ReadTaskRequestsReceived);
-    auto response = (*extension->task_iterator)();
+
+    auto response = (*extension->task_iterator)(extension->replica_info->number_of_current_replica);
     connections->sendReadTaskResponse(response);
 }
 

src/QueryPipeline/RemoteQueryExecutor.h

@@ -28,7 +28,7 @@ class RemoteQueryExecutorReadContext;
 class ParallelReplicasReadingCoordinator;
 
 /// This is the same type as StorageS3Source::IteratorWrapper
-using TaskIterator = std::function<String()>;
+using TaskIterator = std::function<String(size_t)>;
 
 /// This class allows one to launch queries on remote replicas of one shard and get results
 class RemoteQueryExecutor

src/Storages/IStorageCluster.cpp

@@ -112,7 +112,7 @@ void ReadFromCluster::createExtension(const ActionsDAG::Node * predicate)
     if (extension)
         return;
 
-    extension = storage->getTaskIteratorExtension(predicate, context);
+    extension = storage->getTaskIteratorExtension(predicate, context, cluster);
 }
 
 /// The code executes on initiator
@@ -178,8 +178,6 @@ void IStorageCluster::read(
 
 void ReadFromCluster::initializePipeline(QueryPipelineBuilder & pipeline, const BuildQueryPipelineSettings &)
 {
-    createExtension(nullptr);
-
     const Scalars & scalars = context->hasQueryContext() ? context->getQueryContext()->getScalars() : Scalars{};
     const bool add_agg_info = processed_stage == QueryProcessingStage::WithMergeableState;
 
@@ -192,6 +190,10 @@ 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);
 
+    size_t replica_index = 0;
+
+    createExtension(nullptr);
+
     for (const auto & shard_info : cluster->getShardsInfo())
     {
         if (pipes.size() >= max_replicas_to_use)
@@ -209,6 +211,8 @@ void ReadFromCluster::initializePipeline(QueryPipelineBuilder & pipeline, const
         if (try_results.empty())
             continue;
 
+        IConnections::ReplicaInfo replica_info{ .number_of_current_replica = replica_index++ };
+
         auto remote_query_executor = std::make_shared<RemoteQueryExecutor>(
             std::vector<IConnectionPool::Entry>{try_results.front()},
             query_to_send->formatWithSecretsOneLine(),
@@ -218,7 +222,7 @@ void ReadFromCluster::initializePipeline(QueryPipelineBuilder & pipeline, const
             scalars,
             Tables(),
             processed_stage,
-            extension);
+            RemoteQueryExecutor::Extension{.task_iterator = extension->task_iterator, .replica_info = std::move(replica_info)});
 
         remote_query_executor->setLogger(log);
         pipes.emplace_back(std::make_shared<RemoteSource>(

src/Storages/IStorageCluster.h

@@ -33,7 +33,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) 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

@@ -13,6 +13,7 @@
 #include <Storages/ObjectStorage/Utils.h>
 #include <Storages/ObjectStorage/StorageObjectStorageSource.h>
 #include <Storages/extractTableFunctionFromSelectQuery.h>
+#include <Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.h>
 
 
 namespace DB
@@ -144,24 +145,34 @@ void StorageObjectStorageCluster::updateQueryToSendIfNeeded(
 }
 
 RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExtension(
-    const ActionsDAG::Node * predicate, const ContextPtr & local_context) 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);
 
-    auto callback = std::make_shared<std::function<String()>>([iterator]() mutable -> String
+    std::vector<std::string> ids_of_hosts;
+    for (const auto & shard : cluster->getShardsInfo())
     {
-        auto object_info = iterator->next(0);
-        if (!object_info)
-            return "";
+        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 archive_object_info = std::dynamic_pointer_cast<StorageObjectStorageSource::ArchiveIterator::ObjectInfoInArchive>(object_info);
-        if (archive_object_info)
-            return archive_object_info->getPathToArchive();
+    auto callback = std::make_shared<TaskIterator>(
+        [task_distributor](size_t number_of_current_replica) mutable -> String {
+            return task_distributor->getNextTask(number_of_current_replica).value_or("");
+        });
 
-        return object_info->getPath();
-    });
     return RemoteQueryExecutor::Extension{ .task_iterator = std::move(callback) };
 }
 

src/Storages/ObjectStorage/StorageObjectStorageCluster.h

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

src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.cpp

@@ -0,0 +1,182 @@
+#include "StorageObjectStorageStableTaskDistributor.h"
+#include <Common/SipHash.h>
+#include <consistent_hashing.h>
+#include <optional>
+
+namespace DB
+{
+
+StorageObjectStorageStableTaskDistributor::StorageObjectStorageStableTaskDistributor(
+    std::shared_ptr<IObjectIterator> iterator_,
+    std::vector<std::string> ids_of_nodes_)
+    : iterator(std::move(iterator_))
+    , connection_to_files(ids_of_nodes_.size())
+    , ids_of_nodes(ids_of_nodes_)
+    , iterator_exhausted(false)
+{
+}
+
+std::optional<String> StorageObjectStorageStableTaskDistributor::getNextTask(size_t number_of_current_replica)
+{
+    LOG_TRACE(
+        log,
+        "Received a new connection from replica {} looking for a file",
+        number_of_current_replica
+    );
+
+    // 1. Check pre-queued files first
+    if (auto file = getPreQueuedFile(number_of_current_replica))
+        return file;
+
+    // 2. Try to find a matching file from the iterator
+    if (auto file = getMatchingFileFromIterator(number_of_current_replica))
+        return file;
+
+    // 3. Process unprocessed files if iterator is exhausted
+    return getAnyUnprocessedFile(number_of_current_replica);
+}
+
+size_t StorageObjectStorageStableTaskDistributor::getReplicaForFile(const String & file_path)
+{
+    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())
+    {
+        String next_file = files.back();
+        files.pop_back();
+
+        auto it = unprocessed_files.find(next_file);
+        if (it == unprocessed_files.end())
+            continue;
+
+        unprocessed_files.erase(it);
+
+        LOG_TRACE(
+            log,
+            "Assigning pre-queued file {} to replica {}",
+            next_file,
+            number_of_current_replica
+        );
+
+        return next_file;
+    }
+
+    return std::nullopt;
+}
+
+std::optional<String> StorageObjectStorageStableTaskDistributor::getMatchingFileFromIterator(size_t number_of_current_replica)
+{
+    {
+        std::lock_guard lock(mutex);
+        if (iterator_exhausted)
+            return std::nullopt;
+    }
+
+    while (true)
+    {
+        ObjectInfoPtr object_info;
+
+        {
+            std::lock_guard lock(mutex);
+            object_info = iterator->next(0);
+
+            if (!object_info)
+            {
+                iterator_exhausted = true;
+                break;
+            }
+        }
+
+        String file_path;
+
+        auto archive_object_info = std::dynamic_pointer_cast<StorageObjectStorageSource::ArchiveIterator::ObjectInfoInArchive>(object_info);
+        if (archive_object_info)
+        {
+            file_path = archive_object_info->getPathToArchive();
+        }
+        else
+        {
+            file_path = object_info->getPath();
+        }
+
+        size_t file_replica_idx = getReplicaForFile(file_path);
+        if (file_replica_idx == number_of_current_replica)
+        {
+            LOG_TRACE(
+                log,
+                "Found file {} for replica {}",
+                file_path,
+                number_of_current_replica
+            );
+
+            return file_path;
+        }
+
+        // Queue file for its assigned replica
+        {
+            std::lock_guard lock(mutex);
+            unprocessed_files.insert(file_path);
+            connection_to_files[file_replica_idx].push_back(file_path);
+        }
+    }
+
+    return std::nullopt;
+}
+
+std::optional<String> StorageObjectStorageStableTaskDistributor::getAnyUnprocessedFile(size_t number_of_current_replica)
+{
+    std::lock_guard lock(mutex);
+
+    if (!unprocessed_files.empty())
+    {
+        auto it = unprocessed_files.begin();
+        String next_file = *it;
+        unprocessed_files.erase(it);
+
+        LOG_TRACE(
+            log,
+            "Iterator exhausted. Assigning unprocessed file {} to replica {}",
+            next_file,
+            number_of_current_replica
+        );
+
+        return next_file;
+    }
+
+    return std::nullopt;
+}
+
+}

src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.h

@@ -0,0 +1,44 @@
+#pragma once
+
+#include <Client/Connection.h>
+#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>
+#include <memory>
+
+namespace DB
+{
+
+class StorageObjectStorageStableTaskDistributor
+{
+public:
+    StorageObjectStorageStableTaskDistributor(
+        std::shared_ptr<IObjectIterator> iterator_,
+        std::vector<std::string> ids_of_nodes_);
+
+    std::optional<String> getNextTask(size_t number_of_current_replica);
+
+private:
+    size_t getReplicaForFile(const String & file_path);
+    std::optional<String> getPreQueuedFile(size_t number_of_current_replica);
+    std::optional<String> getMatchingFileFromIterator(size_t number_of_current_replica);
+    std::optional<String> getAnyUnprocessedFile(size_t number_of_current_replica);
+
+    std::shared_ptr<IObjectIterator> iterator;
+
+    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;
+
+    LoggerPtr log = getLogger("StorageClusterTaskDistributor");
+};
+
+}