Merge pull request #780 from Altinity/feature/retries_in_cluster_functions

Restart cluster tasks on connection lost

Author: Enmk

src/Core/Protocol.h

@@ -96,8 +96,10 @@ namespace Protocol
             MergeTreeReadTaskRequest = 16,  /// Request from a MergeTree replica to a coordinator
             TimezoneUpdate = 17,            /// Receive server's (session-wide) default timezone
             SSHChallenge = 18,              /// Return challenge for SSH signature signing
+
             MAX = SSHChallenge,
 
+            ConnectionLost = 255,           /// Exception that occurred on the client side.
         };
 
         /// NOTE: If the type of packet argument would be Enum, the comparison packet >= 0 && packet < 10

src/Core/Settings.cpp

@@ -7027,6 +7027,9 @@ DECLARE(Bool, allow_experimental_ytsaurus_dictionary_source, false, R"(
     )", EXPERIMENTAL) \
     DECLARE(Bool, distributed_plan_force_shuffle_aggregation, false, R"(
 Use Shuffle aggregation strategy instead of PartialAggregation + Merge in distributed query plan.
+)", EXPERIMENTAL) \
+    DECLARE(Bool, allow_retries_in_cluster_requests, false, R"(
+Allow retries in cluster request, when one node goes offline
 )", EXPERIMENTAL) \
     \
     /** Experimental timeSeries* aggregate functions. */ \

src/Core/SettingsChangesHistory.cpp

@@ -260,6 +260,8 @@ const VersionToSettingsChangesMap & getSettingsChangesHistory()
             {"parallel_replicas_for_cluster_engines", false, true, "New setting."},
             {"parallel_hash_join_threshold", 0, 0, "New setting"},
             /// Release closed. Please use 25.4
+            {"use_object_storage_list_objects_cache", true, false, "New setting."},
+            {"allow_retries_in_cluster_requests", false, false, "New setting."},
         });
         addSettingsChanges(settings_changes_history, "25.2",
         {

src/QueryPipeline/RemoteQueryExecutor.cpp

@@ -52,6 +52,7 @@ namespace Setting
     extern const SettingsBool use_hedged_requests;
     extern const SettingsBool push_external_roles_in_interserver_queries;
     extern const SettingsMilliseconds parallel_replicas_connect_timeout_ms;
+    extern const SettingsBool allow_retries_in_cluster_requests;
 }
 
 namespace ErrorCodes
@@ -82,6 +83,7 @@ RemoteQueryExecutor::RemoteQueryExecutor(
     , extension(extension_)
     , priority_func(priority_func_)
     , read_packet_type_separately(context->canUseParallelReplicasOnInitiator() && !context->getSettingsRef()[Setting::use_hedged_requests])
+    , allow_retries_in_cluster_requests(context->getSettingsRef()[Setting::allow_retries_in_cluster_requests])
 {
     if (stage == QueryProcessingStage::QueryPlan && !query_plan)
         throw Exception(ErrorCodes::LOGICAL_ERROR, "Query plan is not passed for QueryPlan processing stage");
@@ -457,7 +459,8 @@ int RemoteQueryExecutor::sendQueryAsync()
         read_context = std::make_unique<ReadContext>(
             *this,
             /*suspend_when_query_sent*/ true,
-            read_packet_type_separately);
+            read_packet_type_separately,
+            allow_retries_in_cluster_requests);
 
     /// If query already sent, do nothing. Note that we cannot use sent_query flag here,
     /// because we can still be in process of sending scalars or external tables.
@@ -530,7 +533,8 @@ RemoteQueryExecutor::ReadResult RemoteQueryExecutor::readAsync()
         read_context = std::make_unique<ReadContext>(
             *this,
             /*suspend_when_query_sent*/ false,
-            read_packet_type_separately);
+            read_packet_type_separately,
+            allow_retries_in_cluster_requests);
         recreate_read_context = false;
     }
 
@@ -654,7 +658,11 @@ RemoteQueryExecutor::ReadResult RemoteQueryExecutor::processPacket(Packet packet
             /// We can actually return it, and the first call to RemoteQueryExecutor::read
             /// will return earlier. We should consider doing it.
             if (!packet.block.empty() && (packet.block.rows() > 0))
+            {
+                if (extension && extension->replica_info)
+                    replica_has_processed_data.insert(extension->replica_info->number_of_current_replica);
                 return ReadResult(adaptBlockStructure(packet.block, *header));
+            }
             break;  /// If the block is empty - we will receive other packets before EndOfStream.
 
         case Protocol::Server::Exception:
@@ -716,6 +724,22 @@ RemoteQueryExecutor::ReadResult RemoteQueryExecutor::processPacket(Packet packet
         case Protocol::Server::TimezoneUpdate:
             break;
 
+        case Protocol::Server::ConnectionLost:
+            if (allow_retries_in_cluster_requests)
+            {
+                if (extension && extension->task_iterator && extension->task_iterator->supportRerunTask() && extension->replica_info)
+                {
+                    if (!replica_has_processed_data.contains(extension->replica_info->number_of_current_replica))
+                    {
+                        finished = true;
+                        extension->task_iterator->rescheduleTasksFromReplica(extension->replica_info->number_of_current_replica);
+                        return ReadResult(Block{});
+                    }
+                }
+            }
+            packet.exception->rethrow();
+            break;
+
         default:
             got_unknown_packet_from_replica = true;
             throw Exception(

src/QueryPipeline/RemoteQueryExecutor.h

@@ -31,8 +31,22 @@ class RemoteQueryExecutorReadContext;
 
 class ParallelReplicasReadingCoordinator;
 
-/// This is the same type as StorageS3Source::IteratorWrapper
-using TaskIterator = std::function<ClusterFunctionReadTaskResponsePtr(size_t)>;
+namespace ErrorCodes
+{
+    extern const int NOT_IMPLEMENTED;
+};
+
+class TaskIterator
+{
+public:
+    virtual ~TaskIterator() = default;
+    virtual bool supportRerunTask() const { return false; }
+    virtual void rescheduleTasksFromReplica(size_t /* number_of_current_replica */)
+    {
+        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Method rescheduleTasksFromReplica is not implemented");
+    }
+    virtual ClusterFunctionReadTaskResponsePtr operator()(size_t number_of_current_replica) const = 0;
+};
 
 /// This class allows one to launch queries on remote replicas of one shard and get results
 class RemoteQueryExecutor
@@ -316,6 +330,10 @@ class RemoteQueryExecutor
 
     const bool read_packet_type_separately = false;
 
+    const bool allow_retries_in_cluster_requests = false;
+
+    std::unordered_set<size_t> replica_has_processed_data;
+
     /// Send all scalars to remote servers
     void sendScalars();
 

src/QueryPipeline/RemoteQueryExecutorReadContext.cpp

@@ -19,11 +19,15 @@ namespace ErrorCodes
 }
 
 RemoteQueryExecutorReadContext::RemoteQueryExecutorReadContext(
-    RemoteQueryExecutor & executor_, bool suspend_when_query_sent_, bool read_packet_type_separately_)
+    RemoteQueryExecutor & executor_,
+    bool suspend_when_query_sent_,
+    bool read_packet_type_separately_,
+    bool allow_retries_in_cluster_requests_)
     : AsyncTaskExecutor(std::make_unique<Task>(*this))
     , executor(executor_)
     , suspend_when_query_sent(suspend_when_query_sent_)
     , read_packet_type_separately(read_packet_type_separately_)
+    , allow_retries_in_cluster_requests(allow_retries_in_cluster_requests_)
 {
     if (-1 == pipe2(pipe_fd, O_NONBLOCK))
         throw ErrnoException(ErrorCodes::CANNOT_OPEN_FILE, "Cannot create pipe");
@@ -54,18 +58,29 @@ void RemoteQueryExecutorReadContext::Task::run(AsyncCallback async_callback, Sus
     if (read_context.executor.needToSkipUnavailableShard())
         return;
 
-    while (true)
+    try
     {
-        read_context.has_read_packet_part = PacketPart::None;
-
-        if (read_context.read_packet_type_separately)
+        while (true)
         {
-            read_context.packet.type = read_context.executor.getConnections().receivePacketTypeUnlocked(async_callback);
-            read_context.has_read_packet_part = PacketPart::Type;
+            read_context.has_read_packet_part = PacketPart::None;
+
+            if (read_context.read_packet_type_separately)
+            {
+                read_context.packet.type = read_context.executor.getConnections().receivePacketTypeUnlocked(async_callback);
+                read_context.has_read_packet_part = PacketPart::Type;
+                suspend_callback();
+            }
+            read_context.packet = read_context.executor.getConnections().receivePacketUnlocked(async_callback);
+            read_context.has_read_packet_part = PacketPart::Body;
             suspend_callback();
         }
-        read_context.packet = read_context.executor.getConnections().receivePacketUnlocked(async_callback);
-        read_context.has_read_packet_part = PacketPart::Body;
+    }
+    catch (const Exception &)
+    {
+        if (!read_context.allow_retries_in_cluster_requests)
+            throw;
+        read_context.packet.type = Protocol::Server::ConnectionLost;
+        read_context.packet.exception = std::make_unique<Exception>(getCurrentExceptionMessageAndPattern(true), getCurrentExceptionCode());
         suspend_callback();
     }
 }

src/QueryPipeline/RemoteQueryExecutorReadContext.h

@@ -26,7 +26,10 @@ class RemoteQueryExecutorReadContext : public AsyncTaskExecutor
 {
 public:
     explicit RemoteQueryExecutorReadContext(
-        RemoteQueryExecutor & executor_, bool suspend_when_query_sent_, bool read_packet_type_separately_);
+        RemoteQueryExecutor & executor_,
+        bool suspend_when_query_sent_,
+        bool read_packet_type_separately_,
+        bool allow_retries_in_cluster_requests_);
 
     ~RemoteQueryExecutorReadContext() override;
 
@@ -108,6 +111,7 @@ class RemoteQueryExecutorReadContext : public AsyncTaskExecutor
     bool suspend_when_query_sent = false;
     bool is_query_sent = false;
     const bool read_packet_type_separately = false;
+    const bool allow_retries_in_cluster_requests = false;
 };
 
 }

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -191,6 +191,36 @@ void StorageObjectStorageCluster::updateQueryToSendIfNeeded(
     }
 }
 
+class TaskDistributor : public TaskIterator
+{
+public:
+    TaskDistributor(std::shared_ptr<IObjectIterator> iterator,
+        std::vector<std::string> && ids_of_hosts,
+        bool send_over_whole_archive,
+        uint64_t lock_object_storage_task_distribution_ms,
+        ContextPtr context_
+        )
+        : task_distributor(iterator, std::move(ids_of_hosts), send_over_whole_archive, lock_object_storage_task_distribution_ms)
+        , context(context_) {}
+    ~TaskDistributor() override = default;
+    bool supportRerunTask() const override { return true; }
+    void rescheduleTasksFromReplica(size_t number_of_current_replica) override
+    {
+        task_distributor.rescheduleTasksFromReplica(number_of_current_replica);
+    }
+
+    ClusterFunctionReadTaskResponsePtr operator()(size_t number_of_current_replica) const override
+    {
+        auto task = task_distributor.getNextTask(number_of_current_replica);
+        if (task)
+            return std::make_shared<ClusterFunctionReadTaskResponse>(std::move(task), context);
+        return std::make_shared<ClusterFunctionReadTaskResponse>();
+    }
+
+private:
+    mutable StorageObjectStorageStableTaskDistributor task_distributor;
+    ContextPtr context;
+};
 
 RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExtension(
     const ActionsDAG::Node * predicate,
@@ -238,20 +268,11 @@ RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExten
             lock_object_storage_task_distribution_ms_max
         );
 
-    auto task_distributor = std::make_shared<StorageObjectStorageStableTaskDistributor>(
-        iterator,
+    auto callback = std::make_shared<TaskDistributor>(iterator,
         std::move(ids_of_hosts),
         /* send_over_whole_archive */!local_context->getSettingsRef()[Setting::cluster_function_process_archive_on_multiple_nodes],
-        lock_object_storage_task_distribution_ms);
-
-    auto callback = std::make_shared<TaskIterator>(
-        [task_distributor, local_context](size_t number_of_current_replica) mutable -> ClusterFunctionReadTaskResponsePtr
-        {
-            auto task = task_distributor->getNextTask(number_of_current_replica);
-            if (task)
-                return std::make_shared<ClusterFunctionReadTaskResponse>(std::move(task), local_context);
-            return std::make_shared<ClusterFunctionReadTaskResponse>();
-        });
+        lock_object_storage_task_distribution_ms,
+        local_context);
 
     return RemoteQueryExecutor::Extension{ .task_iterator = std::move(callback) };
 }

src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.cpp

@@ -9,7 +9,8 @@ namespace DB
 namespace ErrorCodes
 {
     extern const int LOGICAL_ERROR;
-}
+    extern const int CANNOT_READ_ALL_DATA;
+};
 
 StorageObjectStorageStableTaskDistributor::StorageObjectStorageStableTaskDistributor(
     std::shared_ptr<IObjectIterator> iterator_,
@@ -23,6 +24,9 @@ StorageObjectStorageStableTaskDistributor::StorageObjectStorageStableTaskDistrib
     , lock_object_storage_task_distribution_us(lock_object_storage_task_distribution_ms_ * 1000)
     , iterator_exhausted(false)
 {
+    size_t nodes = ids_of_nodes.size();
+    for (size_t i = 0; i < nodes; ++i)
+        replica_to_files_to_be_processed[i] = std::list<ObjectInfoPtr>{};
 }
 
 ObjectInfoPtr StorageObjectStorageStableTaskDistributor::getNextTask(size_t number_of_current_replica)
@@ -31,16 +35,27 @@ ObjectInfoPtr StorageObjectStorageStableTaskDistributor::getNextTask(size_t numb
 
     saveLastNodeActivity(number_of_current_replica);
 
-    // 1. Check pre-queued files first
-    if (auto file = getPreQueuedFile(number_of_current_replica))
-        return file;
+    auto processed_file_list_ptr = replica_to_files_to_be_processed.find(number_of_current_replica);
+    if (processed_file_list_ptr == replica_to_files_to_be_processed.end())
+        throw Exception(
+            ErrorCodes::LOGICAL_ERROR,
+            "Replica number {} was marked as lost, can't set task for it anymore",
+            number_of_current_replica
+        );
 
+    // 1. Check pre-queued files first
+    auto file = getPreQueuedFile(number_of_current_replica);
     // 2. Try to find a matching file from the iterator
-    if (auto file = getMatchingFileFromIterator(number_of_current_replica))
-        return file;
-
+    if (!file)
+        file = getMatchingFileFromIterator(number_of_current_replica);
     // 3. Process unprocessed files if iterator is exhausted
-    return getAnyUnprocessedFile(number_of_current_replica);
+    if (!file)
+        file = getAnyUnprocessedFile(number_of_current_replica);
+
+    if (file)
+        processed_file_list_ptr->second.push_back(file);
+
+    return file;
 }
 
 size_t StorageObjectStorageStableTaskDistributor::getReplicaForFile(const String & file_path)
@@ -52,16 +67,27 @@ size_t StorageObjectStorageStableTaskDistributor::getReplicaForFile(const String
         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)
+    auto replica = replica_to_files_to_be_processed.begin();
+    if (replica == replica_to_files_to_be_processed.end())
+        throw Exception(
+            ErrorCodes::LOGICAL_ERROR,
+            "No active replicas, can't find best replica for file {}",
+            file_path
+        );
+
+    size_t best_id = replica->first;
+    UInt64 best_weight = sipHash64(ids_of_nodes[best_id] + file_path);
+    ++replica;
+    while (replica != replica_to_files_to_be_processed.end())
     {
+        size_t id = replica->first;
         UInt64 weight = sipHash64(ids_of_nodes[id] + file_path);
         if (weight > best_weight)
         {
             best_weight = weight;
             best_id = id;
         }
+        ++replica;
     }
     return best_id;
 }
@@ -234,4 +260,28 @@ void StorageObjectStorageStableTaskDistributor::saveLastNodeActivity(size_t numb
     last_node_activity[number_of_current_replica] = now;
 }
 
+void StorageObjectStorageStableTaskDistributor::rescheduleTasksFromReplica(size_t number_of_current_replica)
+{
+    LOG_INFO(log, "Replica {} is marked as lost, tasks are returned to queue", number_of_current_replica);
+    std::lock_guard lock(mutex);
+
+    auto processed_file_list_ptr = replica_to_files_to_be_processed.find(number_of_current_replica);
+    if (processed_file_list_ptr == replica_to_files_to_be_processed.end())
+        throw Exception(
+            ErrorCodes::LOGICAL_ERROR,
+            "Replica number {} was marked as lost already",
+            number_of_current_replica
+        );
+
+    if (replica_to_files_to_be_processed.size() < 2)
+        throw Exception(
+            ErrorCodes::CANNOT_READ_ALL_DATA,
+            "All replicas were marked as lost"
+        );
+
+    replica_to_files_to_be_processed.erase(number_of_current_replica);
+    for (const auto & file : processed_file_list_ptr->second)
+        unprocessed_files.emplace(file->getPath(), std::make_pair(file, getReplicaForFile(file->getPath())));
+}
+
 }