From aa4b79780845a78388571c241645de5ab194d881 Mon Sep 17 00:00:00 2001
From: kssenii <sumarokovakseniia@mail.ru>
Date: Sun, 10 Oct 2021 23:22:58 +0300
Subject: [PATCH] Turn off tasks stealing for remote disk

---
 src/Storages/MergeTree/IMergeTreeDataPart.h         | 2 ++
 src/Storages/MergeTree/MergeTreeDataPartCompact.cpp | 5 +++++
 src/Storages/MergeTree/MergeTreeDataPartCompact.h   | 2 ++
 src/Storages/MergeTree/MergeTreeDataPartInMemory.h  | 1 +
 src/Storages/MergeTree/MergeTreeDataPartWide.cpp    | 5 +++++
 src/Storages/MergeTree/MergeTreeDataPartWide.h      | 2 ++
 src/Storages/MergeTree/MergeTreeReadPool.cpp        | 4 ++++
 7 files changed, 21 insertions(+)

diff --git a/src/Storages/MergeTree/IMergeTreeDataPart.h b/src/Storages/MergeTree/IMergeTreeDataPart.h
index be48aed5c8b..f02b7e98f02 100644
--- a/src/Storages/MergeTree/IMergeTreeDataPart.h
+++ b/src/Storages/MergeTree/IMergeTreeDataPart.h
@@ -95,6 +95,8 @@ public:
 
     virtual bool isStoredOnDisk() const = 0;
 
+    virtual bool isStoredOnRemoteDisk() const = 0;
+
     virtual bool supportsVerticalMerge() const { return false; }
 
     /// NOTE: Returns zeros if column files are not found in checksums.
diff --git a/src/Storages/MergeTree/MergeTreeDataPartCompact.cpp b/src/Storages/MergeTree/MergeTreeDataPartCompact.cpp
index ad1c2abeee7..c4c2e65547b 100644
--- a/src/Storages/MergeTree/MergeTreeDataPartCompact.cpp
+++ b/src/Storages/MergeTree/MergeTreeDataPartCompact.cpp
@@ -182,6 +182,11 @@ void MergeTreeDataPartCompact::checkConsistency(bool require_part_metadata) cons
     }
 }
 
+bool MergeTreeDataPartCompact::isStoredOnRemoteDisk() const
+{
+    return volume->getDisk()->isRemote();
+}
+
 MergeTreeDataPartCompact::~MergeTreeDataPartCompact()
 {
     removeIfNeeded();
diff --git a/src/Storages/MergeTree/MergeTreeDataPartCompact.h b/src/Storages/MergeTree/MergeTreeDataPartCompact.h
index 2fcc7b7034b..38bfa11652a 100644
--- a/src/Storages/MergeTree/MergeTreeDataPartCompact.h
+++ b/src/Storages/MergeTree/MergeTreeDataPartCompact.h
@@ -56,6 +56,8 @@ public:
 
     bool isStoredOnDisk() const override { return true; }
 
+    bool isStoredOnRemoteDisk() const override;
+
     bool hasColumnFiles(const NameAndTypePair & column) const override;
 
     String getFileNameForColumn(const NameAndTypePair & /* column */) const override { return DATA_FILE_NAME; }
diff --git a/src/Storages/MergeTree/MergeTreeDataPartInMemory.h b/src/Storages/MergeTree/MergeTreeDataPartInMemory.h
index 118340f0233..d1a0344859d 100644
--- a/src/Storages/MergeTree/MergeTreeDataPartInMemory.h
+++ b/src/Storages/MergeTree/MergeTreeDataPartInMemory.h
@@ -44,6 +44,7 @@ public:
         const MergeTreeIndexGranularity & computed_index_granularity) const override;
 
     bool isStoredOnDisk() const override { return false; }
+    bool isStoredOnRemoteDisk() const override { return false; }
     bool hasColumnFiles(const NameAndTypePair & column) const override { return !!getColumnPosition(column.name); }
     String getFileNameForColumn(const NameAndTypePair & /* column */) const override { return ""; }
     void renameTo(const String & new_relative_path, bool remove_new_dir_if_exists) const override;
diff --git a/src/Storages/MergeTree/MergeTreeDataPartWide.cpp b/src/Storages/MergeTree/MergeTreeDataPartWide.cpp
index 11e080fda6c..6ac4919329d 100644
--- a/src/Storages/MergeTree/MergeTreeDataPartWide.cpp
+++ b/src/Storages/MergeTree/MergeTreeDataPartWide.cpp
@@ -142,6 +142,11 @@ void MergeTreeDataPartWide::loadIndexGranularity()
     index_granularity.setInitialized();
 }
 
+bool MergeTreeDataPartWide::isStoredOnRemoteDisk() const
+{
+    return volume->getDisk()->isRemote();
+}
+
 MergeTreeDataPartWide::~MergeTreeDataPartWide()
 {
     removeIfNeeded();
diff --git a/src/Storages/MergeTree/MergeTreeDataPartWide.h b/src/Storages/MergeTree/MergeTreeDataPartWide.h
index a43396f8cd5..4796143e11d 100644
--- a/src/Storages/MergeTree/MergeTreeDataPartWide.h
+++ b/src/Storages/MergeTree/MergeTreeDataPartWide.h
@@ -50,6 +50,8 @@ public:
 
     bool isStoredOnDisk() const override { return true; }
 
+    bool isStoredOnRemoteDisk() const override;
+
     bool supportsVerticalMerge() const override { return true; }
 
     String getFileNameForColumn(const NameAndTypePair & column) const override;
diff --git a/src/Storages/MergeTree/MergeTreeReadPool.cpp b/src/Storages/MergeTree/MergeTreeReadPool.cpp
index d08cec24184..f23c58bc2b3 100644
--- a/src/Storages/MergeTree/MergeTreeReadPool.cpp
+++ b/src/Storages/MergeTree/MergeTreeReadPool.cpp
@@ -221,6 +221,10 @@ std::vector<size_t> MergeTreeReadPool::fillPerPartInfo(
     {
         const auto & part = parts[i];
 
+        /// Turn off tasks stealing in case there is remote disk.
+        if (part.data_part->isStoredOnRemoteDisk())
+            do_not_steal_tasks = true;
+
         /// Read marks for every data part.
         size_t sum_marks = 0;
         for (const auto & range : part.ranges)