Restart cluster tasks on connection lost

Author: ianton-ru

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/QueryPipeline/RemoteQueryExecutor.cpp

@@ -716,6 +716,16 @@ RemoteQueryExecutor::ReadResult RemoteQueryExecutor::processPacket(Packet packet
         case Protocol::Server::TimezoneUpdate:
             break;
 
+        case Protocol::Server::ConnectionLost:
+            if (extension && extension->task_iterator && extension->task_iterator->supportRerunTask() && extension->replica_info)
+            {
+                finished = true;
+                extension->task_iterator->rerunTasksForReplica(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

@@ -8,6 +8,8 @@
 #include <Interpreters/StorageID.h>
 #include <sys/types.h>
 
+#include <list>
+
 
 namespace DB
 {
@@ -27,8 +29,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 rerunTasksForReplica(size_t /* number_of_current_replica */)
+    {
+        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Method rerunTasksForReplica 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

src/QueryPipeline/RemoteQueryExecutorReadContext.cpp

@@ -54,18 +54,27 @@ 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 &)
+    {
+        read_context.packet.type = Protocol::Server::ConnectionLost;
+        read_context.packet.exception = std::make_unique<Exception>(getCurrentExceptionMessageAndPattern(true), getCurrentExceptionCode());
         suspend_callback();
     }
 }

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -337,6 +337,28 @@ 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)
+        : task_distributor(iterator, ids_of_hosts) {}
+    ~TaskDistributor() override = default;
+    bool supportRerunTask() const override { return true; }
+    void rerunTasksForReplica(size_t number_of_current_replica) override
+    {
+        task_distributor.rerunTasksForReplica(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,
@@ -360,12 +382,7 @@ RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExten
         }
     }
 
-    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 {
-            return task_distributor->getNextTask(number_of_current_replica).value_or("");
-        });
+    auto callback = std::make_shared<TaskDistributor>(iterator, ids_of_hosts);
 
     return RemoteQueryExecutor::Extension{ .task_iterator = std::move(callback) };
 }

src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.cpp

@@ -6,6 +6,12 @@
 namespace DB
 {
 
+namespace ErrorCodes
+{
+    extern const int LOGICAL_ERROR;
+    extern const int CANNOT_READ_ALL_DATA;
+};
+
 StorageObjectStorageStableTaskDistributor::StorageObjectStorageStableTaskDistributor(
     std::shared_ptr<IObjectIterator> iterator_,
     std::vector<std::string> ids_of_nodes_)
@@ -14,6 +20,9 @@ StorageObjectStorageStableTaskDistributor::StorageObjectStorageStableTaskDistrib
     , ids_of_nodes(ids_of_nodes_)
     , iterator_exhausted(false)
 {
+    size_t nodes = ids_of_nodes.size();
+    for (size_t i = 0; i < nodes; ++i)
+        processed_files[i] = std::list<String>{};
 }
 
 std::optional<String> StorageObjectStorageStableTaskDistributor::getNextTask(size_t number_of_current_replica)
@@ -24,16 +33,27 @@ std::optional<String> StorageObjectStorageStableTaskDistributor::getNextTask(siz
         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 = processed_files.find(number_of_current_replica);
+    if (processed_file_list_ptr == processed_files.end())
+        throw Exception(
+            ErrorCodes::LOGICAL_ERROR,
+            "Replica number {} was marked as lost, can't set satk for it anymore",
+            number_of_current_replica
+        );
 
+    // 1. Check pre-queued files first
+    std::optional<String> 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.has_value())
+        file = getMatchingFileFromIterator(number_of_current_replica);
     // 3. Process unprocessed files if iterator is exhausted
-    return getAnyUnprocessedFile(number_of_current_replica);
+    if (!file.has_value())
+        file = getAnyUnprocessedFile(number_of_current_replica);
+
+    if (file.has_value())
+        processed_file_list_ptr->second.push_back(*file);
+
+    return file;
 }
 
 size_t StorageObjectStorageStableTaskDistributor::getReplicaForFile(const String & file_path)
@@ -179,4 +199,28 @@ std::optional<String> StorageObjectStorageStableTaskDistributor::getAnyUnprocess
     return std::nullopt;
 }
 
+void StorageObjectStorageStableTaskDistributor::rerunTasksForReplica(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 = processed_files.find(number_of_current_replica);
+    if (processed_file_list_ptr == processed_files.end())
+        throw Exception(
+            ErrorCodes::LOGICAL_ERROR,
+            "Replica number {} was marked as lost already",
+            number_of_current_replica
+        );
+
+    if (processed_files.size() < 2)
+        throw Exception(
+            ErrorCodes::CANNOT_READ_ALL_DATA,
+            "All replicas were marked as lost"
+        );
+
+    for (const auto & file_path : processed_file_list_ptr->second)
+        unprocessed_files.insert(file_path);
+    processed_files.erase(number_of_current_replica);
+}
+
 }

src/Storages/ObjectStorage/StorageObjectStorageStableTaskDistributor.h

@@ -6,6 +6,8 @@
 #include <Storages/ObjectStorage/StorageObjectStorageSource.h>
 #include <Storages/ObjectStorageQueue/ObjectStorageQueueSource.h>
 #include <unordered_set>
+#include <unordered_map>
+#include <list>
 #include <vector>
 #include <mutex>
 #include <memory>
@@ -22,6 +24,9 @@ class StorageObjectStorageStableTaskDistributor
 
     std::optional<String> getNextTask(size_t number_of_current_replica);
 
+    /// Insert objects back to unprocessed files
+    void rerunTasksForReplica(size_t number_of_current_replica);
+
 private:
     size_t getReplicaForFile(const String & file_path);
     std::optional<String> getPreQueuedFile(size_t number_of_current_replica);
@@ -34,6 +39,7 @@ class StorageObjectStorageStableTaskDistributor
     std::unordered_set<String> unprocessed_files;
 
     std::vector<std::string> ids_of_nodes;
+    std::unordered_map<size_t, std::list<String>> processed_files;
 
     std::mutex mutex;
     bool iterator_exhausted = false;

src/Storages/StorageFileCluster.cpp

@@ -74,14 +74,35 @@ void StorageFileCluster::updateQueryToSendIfNeeded(DB::ASTPtr & query, const Sto
         expression_list->children, storage_snapshot->metadata->getColumns().getAll().toNamesAndTypesDescription(), format_name, context);
 }
 
+class FileTaskIterator : public TaskIterator
+{
+public:
+    FileTaskIterator(const Strings & files,
+        std::optional<StorageFile::ArchiveInfo> archive_info,
+        const ActionsDAG::Node * predicate,
+        const NamesAndTypesList & virtual_columns,
+        const ContextPtr & context,
+        bool distributed_processing = false)
+        : iterator(files, archive_info, predicate, virtual_columns, context, distributed_processing) {}
+
+    ~FileTaskIterator() override = default;
+
+    std::string operator()(size_t /* number_of_current_replica */) const override
+    {
+        return iterator.next();
+    }
+
+private:
+    mutable StorageFileSource::FilesIterator iterator;
+};
+
 RemoteQueryExecutor::Extension StorageFileCluster::getTaskIteratorExtension(
     const ActionsDAG::Node * predicate,
     const std::optional<ActionsDAG> & /* filter_actions_dag */,
     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(); });
+    auto callback = std::make_shared<FileTaskIterator>(paths, std::nullopt, predicate, getVirtualsList(), context);
     return RemoteQueryExecutor::Extension{.task_iterator = std::move(callback)};
 }
 

src/Storages/StorageURLCluster.cpp

@@ -93,15 +93,35 @@ void StorageURLCluster::updateQueryToSendIfNeeded(ASTPtr & query, const StorageS
         expression_list->children, storage_snapshot->metadata->getColumns().getAll().toNamesAndTypesDescription(), format_name, context);
 }
 
+class UrlTaskIterator : public TaskIterator
+{
+public:
+    UrlTaskIterator(const String & uri,
+        size_t max_addresses,
+        const ActionsDAG::Node * predicate,
+        const NamesAndTypesList & virtual_columns,
+        const ContextPtr & context)
+        : iterator(uri, max_addresses, predicate, virtual_columns, context) {}
+
+    ~UrlTaskIterator() override = default;
+
+    std::string operator()(size_t /* number_of_current_replica */) const override
+    {
+        return iterator.next();
+    }
+
+private:
+    mutable StorageURLSource::DisclosedGlobIterator iterator;
+};
+
 RemoteQueryExecutor::Extension StorageURLCluster::getTaskIteratorExtension(
     const ActionsDAG::Node * predicate,
     const std::optional<ActionsDAG> & /* filter_actions_dag */,
     const ContextPtr & context,
     ClusterPtr) const
 {
-    auto iterator = std::make_shared<StorageURLSource::DisclosedGlobIterator>(
+    auto callback = std::make_shared<UrlTaskIterator>(
         uri, context->getSettingsRef()[Setting::glob_expansion_max_elements], predicate, getVirtualsList(), context);
-    auto callback = std::make_shared<TaskIterator>([iter = std::move(iterator)](size_t) mutable -> String { return iter->next(); });
     return RemoteQueryExecutor::Extension{.task_iterator = std::move(callback)};
 }