Merge branch 'antalya-25.6.5' into feature/system_preshutdown

Enmk · web-flow · commit 2d05f2af9fd7 · 2025-09-11T16:57:37.000+02:00
diff --git a/src/Core/Protocol.h b/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
diff --git a/src/Core/Settings.cpp b/src/Core/Settings.cpp
@@ -6883,6 +6883,9 @@ Possible values:
  - '' - do not force any kind of Exchange operators, let the optimizer choose,
  - 'Persisted' - use temporary files in object storage,
  - 'Streaming' - stream exchange data over network.
+)", EXPERIMENTAL) \
+    DECLARE(Bool, allow_retries_in_cluster_requests, false, R"(
+Allow retries in cluster request, when one node goes offline
 )", EXPERIMENTAL) \
     DECLARE(Bool, object_storage_remote_initiator, false, R"(
 Execute request to object storage as remote on one of object_storage_cluster nodes.
diff --git a/src/Core/SettingsChangesHistory.cpp b/src/Core/SettingsChangesHistory.cpp
@@ -196,6 +196,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, "24.12.2.20000",
         {
diff --git a/src/QueryPipeline/RemoteQueryExecutor.cpp b/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");
@@ -484,7 +486,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.
@@ -557,7 +560,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;
     }
 
@@ -681,7 +685,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 && (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:
@@ -743,6 +751,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(
diff --git a/src/QueryPipeline/RemoteQueryExecutor.h b/src/QueryPipeline/RemoteQueryExecutor.h
@@ -28,8 +28,22 @@ class RemoteQueryExecutorReadContext;
 
 class ParallelReplicasReadingCoordinator;
 
-/// This is the same type as StorageS3Source::IteratorWrapper
-using TaskIterator = std::function<String(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 std::string 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
@@ -333,6 +347,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();
 
diff --git a/src/QueryPipeline/RemoteQueryExecutorReadContext.cpp b/src/QueryPipeline/RemoteQueryExecutorReadContext.cpp
@@ -20,11 +20,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");
@@ -55,37 +59,46 @@ void RemoteQueryExecutorReadContext::Task::run(AsyncCallback async_callback, Sus
     if (read_context.executor.needToSkipUnavailableShard())
         return;
 
-    while (true)
+    try
     {
-        try
+        while (true)
         {
-            read_context.has_read_packet_part = PacketPart::None;
-
-            if (read_context.read_packet_type_separately)
+            try
             {
-                read_context.packet.type = read_context.executor.getConnections().receivePacketTypeUnlocked(async_callback);
-                read_context.has_read_packet_part = PacketPart::Type;
-                suspend_callback();
+                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;
+                if (read_context.packet.type == Protocol::Server::Data)
+                    read_context.has_data_packets = true;
             }
-            read_context.packet = read_context.executor.getConnections().receivePacketUnlocked(async_callback);
-            read_context.has_read_packet_part = PacketPart::Body;
-            if (read_context.packet.type == Protocol::Server::Data)
-                read_context.has_data_packets = true;
-        }
-        catch (const Exception & e)
-        {
-            /// If cluster node unxepectedly shutted down (kill/segfault/power off/etc.) socket just closes.
-            /// If initiator did not process any data packets before, this fact can be ignored.
-            /// Unprocessed tasks will be executed on other nodes.
-            if (e.code() == ErrorCodes::ATTEMPT_TO_READ_AFTER_EOF
-                && !read_context.has_data_packets.load() && read_context.executor.skipUnavailableShards())
+            catch (const Exception & e)
             {
-                read_context.has_read_packet_part = PacketPart::None;
+                /// If cluster node unxepectedly shutted down (kill/segfault/power off/etc.) socket just closes.
+                /// If initiator did not process any data packets before, this fact can be ignored.
+                /// Unprocessed tasks will be executed on other nodes.
+                if (e.code() == ErrorCodes::ATTEMPT_TO_READ_AFTER_EOF
+                    && !read_context.has_data_packets.load() && read_context.executor.skipUnavailableShards())
+                {
+                    read_context.has_read_packet_part = PacketPart::None;
+                }
+                else
+                    throw;
             }
-            else
-                throw;
         }
-
+    }
+    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();
     }
 }
diff --git a/src/QueryPipeline/RemoteQueryExecutorReadContext.h b/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;
 
@@ -109,6 +112,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;
 };
 
 }
diff --git a/src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp b/src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp
@@ -439,6 +439,30 @@ void StorageObjectStorageCluster::updateQueryToSendIfNeeded(
         args.insert(args.end(), object_storage_type_arg);
 }
 
+class TaskDistributor : public TaskIterator
+{
+public:
+    TaskDistributor(std::shared_ptr<IObjectIterator> iterator,
+        const std::vector<std::string> & ids_of_hosts,
+        uint64_t lock_object_storage_task_distribution_ms
+        )
+        : task_distributor(iterator, ids_of_hosts, lock_object_storage_task_distribution_ms) {}
+    ~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);
+    }
+
+    std::string operator()(size_t number_of_current_replica) const override
+    {
+        return task_distributor.getNextTask(number_of_current_replica).value_or("");
+    }
+
+private:
+    mutable StorageObjectStorageStableTaskDistributor task_distributor;
+};
+
 RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExtension(
     const ActionsDAG::Node * predicate,
     const std::optional<ActionsDAG> & filter_actions_dag,
@@ -474,14 +498,7 @@ RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExten
             lock_object_storage_task_distribution_ms_max
         );
 
-    auto task_distributor = std::make_shared<StorageObjectStorageStableTaskDistributor>(
-        iterator,
-        ids_of_hosts,
-        lock_object_storage_task_distribution_ms);
-
-    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(""); });
+    auto callback = std::make_shared<TaskDistributor>(iterator, ids_of_hosts, lock_object_storage_task_distribution_ms);
 
     return RemoteQueryExecutor::Extension{.task_iterator = std::move(callback)};
 }
diff --git a/src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.cpp b/src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.cpp
diff --git a/src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.h b/src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.h
diff --git a/src/Storages/StorageFileCluster.cpp b/src/Storages/StorageFileCluster.cpp
diff --git a/src/Storages/StorageURLCluster.cpp b/src/Storages/StorageURLCluster.cpp

Original file line number	Diff line number	Diff line change
`@@ -196,6 +196,8 @@ const VersionToSettingsChangesMap & getSettingsChangesHistory()`
`196`	`196`	`{"parallel_replicas_for_cluster_engines", false, true, "New setting."},`
`197`	`197`	`{"parallel_hash_join_threshold", 0, 0, "New setting"},`
`198`	`198`	`/// Release closed. Please use 25.4`
	`199`	`+ {"use_object_storage_list_objects_cache", true, false, "New setting."},`
	`200`	`+ {"allow_retries_in_cluster_requests", false, false, "New setting."},`
`199`	`201`	`});`
`200`	`202`	`addSettingsChanges(settings_changes_history, "24.12.2.20000",`
`201`	`203`	`{`