Enable manually adding some malfunctioned WN into blocklist (pingcap#308)

Signed-off-by: Calvin Neo <[email protected]>

Author: CalvinNeo

Diff for: contrib/client-c

@@ -192,6 +192,8 @@ struct ContextShared
 
     JointThreadInfoJeallocMapPtr joint_memory_allocation_map; /// Joint thread-wise alloc/dealloc map
 
+    std::unordered_set<uint64_t> store_id_blocklist; /// Those store id are blocked from batch cop request.
+
     class SessionKeyHash
     {
     public:
@@ -2264,6 +2266,48 @@ void Context::setMockMPPServerInfo(MockMPPServerInfo & info)
     mpp_server_info = info;
 }
 
+const std::unordered_set<uint64_t> * Context::getStoreIdBlockList() const
+{
+    return &shared->store_id_blocklist;
+}
+
+// NOLINTNEXTLINE(readability-convert-member-functions-to-static)
+bool Context::initializeStoreIdBlockList(const String & comma_sep_string)
+{
+#if SERVERLESS_PROXY == 1
+    std::istringstream iss(comma_sep_string);
+    std::string token;
+
+    while (std::getline(iss, token, ','))
+    {
+        try
+        {
+            uint64_t number = std::stoull(token);
+            shared->store_id_blocklist.insert(number);
+        }
+        catch (...)
+        {
+            // Keep empty
+            LOG_INFO(DB::Logger::get(), "Error disagg_blocklist_wn_store_id setting, {}", comma_sep_string);
+            shared->store_id_blocklist.clear();
+            return false;
+        }
+    }
+
+    if (!shared->store_id_blocklist.empty())
+        LOG_DEBUG(
+            DB::Logger::get(),
+            "Blocklisted {} stores, which are {}",
+            shared->store_id_blocklist.size(),
+            comma_sep_string);
+
+    return true;
+#else
+    UNUSED(comma_sep_string);
+    return true;
+#endif
+}
+
 SessionCleaner::~SessionCleaner()
 {
     try
@@ -2297,4 +2341,5 @@ void SessionCleaner::run()
             break;
     }
 }
+
 } // namespace DB

Diff for: dbms/src/Interpreters/Context.cpp

@@ -569,6 +569,9 @@ class Context
     bool isRegionInBlacklist(const RegionID region_id);
     bool isRegionsContainsInBlacklist(const std::vector<RegionID> regions);
 
+    bool initializeStoreIdBlockList(const String &);
+    const std::unordered_set<uint64_t> * getStoreIdBlockList() const;
+
 private:
     /** Check if the current client has access to the specified database.
       * If access is denied, throw an exception.

Diff for: dbms/src/Interpreters/Context.h

@@ -55,6 +55,7 @@ struct Settings
     M(SettingUInt64, disagg_build_task_timeout, DEFAULT_DISAGG_TASK_BUILD_TIMEOUT_SEC, "disagg task establish timeout, unit is second.")                                                                                                \
     M(SettingUInt64, disagg_task_snapshot_timeout, DEFAULT_DISAGG_TASK_TIMEOUT_SEC, "disagg task snapshot max endurable time, unit is second.")                                                                                         \
     M(SettingUInt64, disagg_fetch_pages_timeout, DEFAULT_DISAGG_FETCH_PAGES_TIMEOUT_SEC, "fetch disagg pages timeout for one segment, unit is second.")                                                                                 \
+    M(SettingString, disagg_blocklist_wn_store_id, "", "comma seperated unsigned integers representing `store_id`s of stores that are blacklisted.")                                                                                    \
     M(SettingInt64, safe_point_update_interval_seconds, 1, "The interval in seconds to update safe point from PD.")                                                                                                                     \
     M(SettingUInt64, min_compress_block_size, DEFAULT_MIN_COMPRESS_BLOCK_SIZE, "The actual size of the block to compress, if the uncompressed data less than max_compress_block_size is no less than this value "                       \
                                                                                "and no less than the volume of data for one mark.")                                                                                                     \
@@ -251,9 +252,8 @@ struct Settings
     M(SettingDouble, remote_gc_ratio, 0.5, "The files with valid rate less than this threshold will be compacted")                                                                                                                      \
     M(SettingInt64, remote_gc_small_size, 128 * 1024, "The files with total size less than this threshold will be compacted")                                                                                                           \
     M(SettingDouble, disagg_read_concurrency_scale, 20.0, "Scale * logical cpu cores = disaggregated read IO concurrency.")                                                                                                             \
-    \
     M(SettingInt64, fap_wait_checkpoint_timeout_seconds, 80, "The max time wait for a usable checkpoint for FAP")                                                                                                                       \
-    M(SettingUInt64, fap_task_timeout_seconds, 120, "The max time FAP can take before fallback")                                                                                                                       \
+    M(SettingUInt64, fap_task_timeout_seconds, 120, "The max time FAP can take before fallback")                                                                                                                                        \
     M(SettingUInt64, fap_handle_concurrency, 25, "The number of threads for handling FAP tasks")                                                                                                                                        \
     \
     M(SettingUInt64, max_rows_in_set, 0, "Maximum size of the set (in number of elements) resulting from the execution of the IN section.")                                                                                             \

Diff for: dbms/src/Interpreters/Settings.h

@@ -842,6 +842,7 @@ struct SettingString
     void write(WriteBuffer & buf) const { writeBinary(value, buf); }
 
     String get() const { return value; }
+    const String & getRef() const { return value; }
 
 private:
     String value;

Diff for: dbms/src/Interpreters/SettingsCommon.h

@@ -1560,6 +1560,9 @@ int Server::main(const std::vector<std::string> & /*args*/)
         }
     }
 
+    // We do not support blocking store by id in OP mode currently.
+    global_context->initializeStoreIdBlockList(global_context->getSettingsRef().disagg_blocklist_wn_store_id);
+
     LOG_INFO(log, "Loading metadata.");
     loadMetadataSystem(*global_context); // Load "system" database. Its engine keeps as Ordinary.
     /// After attaching system databases we can initialize system log.

Diff for: dbms/src/Server/Server.cpp

@@ -179,6 +179,7 @@ std::vector<pingcap::coprocessor::BatchCopTask> StorageDisaggregated::buildBatch
         table_scan.isPartitionTableScan(),
         physical_table_ids,
         ranges_for_each_physical_table,
+        context.getStoreIdBlockList(),
         store_type,
         label_filter,
         &Poco::Logger::get("pingcap/coprocessor"),