[BugFix] Fix check bind_cpus for resource group (StarRocks#50809)

Signed-off-by: zihe.liu <[email protected]>

be/src/exec/workgroup/pipeline_executor_set.cpp

@@ -26,6 +26,16 @@ namespace starrocks::workgroup {
 // PipelineExecutorSetConfig
 // ------------------------------------------------------------------------------------
 
+static CpuUtil::CpuIds limit_total_cpuids(CpuUtil::CpuIds&& total_cpuids, uint32_t num_total_cores) {
+ if (total_cpuids.empty() || total_cpuids.size() <= num_total_cores) {
+ return std::move(total_cpuids);
+ }
+
+ CpuUtil::CpuIds cpuids;
+ std::copy_n(total_cpuids.begin(), num_total_cores, std::back_inserter(cpuids));
+ return cpuids;
+}
+
 PipelineExecutorSetConfig::PipelineExecutorSetConfig(uint32_t num_total_cores, uint32_t num_total_driver_threads,
  uint32_t num_total_scan_threads,
  uint32_t num_total_connector_scan_threads,
@@ -35,9 +45,9 @@ PipelineExecutorSetConfig::PipelineExecutorSetConfig(uint32_t num_total_cores, u
  num_total_driver_threads(num_total_driver_threads),
  num_total_scan_threads(num_total_scan_threads),
  num_total_connector_scan_threads(num_total_connector_scan_threads),
- total_cpuids(std::move(total_cpuids)),
+ total_cpuids(limit_total_cpuids(std::move(total_cpuids), num_total_cores)),
  enable_bind_cpus(enable_bind_cpus),
- enable_cpu_borrowing(enable_cpu_borrowing) {}
+ enable_cpu_borrowing(enable_cpu_borrowing && enable_bind_cpus) {}
 
 std::string PipelineExecutorSetConfig::to_string() const {
  return fmt::format(
@@ -172,7 +182,7 @@ void PipelineExecutorSet::notify_config_changed() const {
 }
 
 uint32_t PipelineExecutorSet::calculate_num_threads(uint32_t num_total_threads) const {
- if (!_borrowed_cpu_ids.empty() || _cpuids.empty()) {
+ if (!_borrowed_cpu_ids.empty()) {
  return num_total_threads;
  }
  return std::max<uint32_t>(1, num_total_threads * _cpuids.size() / _conf.num_total_cores);

be/src/exec/workgroup/pipeline_executor_set.h

@@ -33,9 +33,9 @@ struct PipelineExecutorSetConfig {
  const uint32_t num_total_scan_threads;
  uint32_t num_total_connector_scan_threads;
 
- CpuUtil::CpuIds total_cpuids;
+ const CpuUtil::CpuIds total_cpuids;
 
- bool enable_bind_cpus;
+ const bool enable_bind_cpus;
  bool enable_cpu_borrowing;
 };
 

be/src/exec/workgroup/pipeline_executor_set_manager.cpp

@@ -147,10 +147,11 @@ void ExecutorsManager::change_num_connector_scan_threads(uint32_t num_connector_
 }
 
 void ExecutorsManager::change_enable_resource_group_cpu_borrowing(bool val) {
- if (_conf.enable_cpu_borrowing == val) {
+ const bool new_val = val && _conf.enable_bind_cpus;
+ if (_conf.enable_cpu_borrowing == new_val) {
  return;
  }
- _conf.enable_cpu_borrowing = val;
+ _conf.enable_cpu_borrowing = new_val;
 
  update_shared_executors();
 }

be/src/http/action/pipeline_blocking_drivers_action.cpp

@@ -161,11 +161,8 @@ void PipelineBlockingDriversAction::_handle_stat(HttpRequest* req) {
  };
 
  QueryMap query_map_in_wg;
- _exec_env->workgroup_manager()->for_each_workgroup([&](const workgroup::WorkGroup& wg) {
- if (wg.exclusive_executors() != nullptr) {
- wg.exclusive_executors()->driver_executor()->iterate_immutable_blocking_driver(
- iterate_func_generator(query_map_in_wg));
- }
+ _exec_env->workgroup_manager()->for_each_executors([&](const workgroup::PipelineExecutorSet& executor) {
+ executor.driver_executor()->iterate_immutable_blocking_driver(iterate_func_generator(query_map_in_wg));
  });
  rapidjson::Document queries_in_wg_obj = query_map_to_doc_func(query_map_in_wg);
 

be/src/runtime/exec_env.cpp

@@ -436,6 +436,7 @@ Status ExecEnv::init(const std::vector<StorePath>& store_paths, bool as_cn) {
  // Disable bind cpus when cgroup has cpu quota but no cpuset.
  const bool enable_bind_cpus = config::enable_resource_group_bind_cpus &&
  (!CpuInfo::is_cgroup_with_cpu_quota() || CpuInfo::is_cgroup_with_cpuset());
+ config::enable_resource_group_bind_cpus = enable_bind_cpus;
  workgroup::PipelineExecutorSetConfig executors_manager_opts(
  CpuInfo::num_cores(), _max_executor_threads, num_io_threads, connector_num_io_threads,
  CpuInfo::get_core_ids(), enable_bind_cpus, config::enable_resource_group_cpu_borrowing);

be/test/CMakeLists.txt

@@ -44,6 +44,7 @@ set(EXEC_FILES
  ./exec/iceberg/iceberg_table_sink_operator_test.cpp
  ./exec/paimon/paimon_delete_file_builder_test.cpp
  ./exec/workgroup/scan_task_queue_test.cpp
+ ./exec/workgroup/pipeline_executor_set_test.cpp
  ./exec/pipeline/pipeline_control_flow_test.cpp
  ./exec/pipeline/pipeline_driver_queue_test.cpp
  ./exec/pipeline/pipeline_file_scan_node_test.cpp

be/test/exec/workgroup/pipeline_executor_set_test.cpp

@@ -0,0 +1,80 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "exec/workgroup/pipeline_executor_set.h"
+
+#include <gtest/gtest.h>
+
+#include "testutil/assert.h"
+#include "testutil/parallel_test.h"
+
+namespace starrocks::workgroup {
+
+PARALLEL_TEST(PipelineExecutorSetConfigTest, test_constructor) {
+ CpuUtil::CpuIds empty_cpuids;
+ CpuUtil::CpuIds cpuids{0, 1, 2, 3, 4, 5, 6, 7};
+
+ /// check enable_cpu_borrowing
+ {
+ PipelineExecutorSetConfig config(2, 2, 2, 2, empty_cpuids, false, true);
+ ASSERT_FALSE(config.enable_cpu_borrowing);
+ }
+
+ {
+ PipelineExecutorSetConfig config(2, 2, 2, 2, empty_cpuids, true, false);
+ ASSERT_FALSE(config.enable_cpu_borrowing);
+ }
+
+ {
+ PipelineExecutorSetConfig config(2, 2, 2, 2, empty_cpuids, true, true);
+ ASSERT_TRUE(config.enable_cpu_borrowing);
+ }
+
+ /// check cpuids
+ {
+ PipelineExecutorSetConfig config(2, 2, 2, 2, empty_cpuids, false, true);
+ ASSERT_EQ(0, config.total_cpuids.size());
+ }
+
+ {
+ PipelineExecutorSetConfig config(0, 2, 2, 2, cpuids, false, true);
+ ASSERT_EQ(0, config.total_cpuids.size());
+ }
+
+ {
+ PipelineExecutorSetConfig config(2, 2, 2, 2, cpuids, false, true);
+ ASSERT_EQ(2, config.total_cpuids.size());
+ for (int i = 0; i < 2; i++) {
+ ASSERT_EQ(i, config.total_cpuids[i]);
+ }
+ }
+
+ {
+ PipelineExecutorSetConfig config(8, 2, 2, 2, cpuids, false, true);
+ ASSERT_EQ(8, config.total_cpuids.size());
+ for (int i = 0; i < 8; i++) {
+ ASSERT_EQ(i, config.total_cpuids[i]);
+ }
+ }
+
+ {
+ PipelineExecutorSetConfig config(100, 2, 2, 2, cpuids, false, true);
+ ASSERT_EQ(8, config.total_cpuids.size());
+ for (int i = 0; i < 8; i++) {
+ ASSERT_EQ(i, config.total_cpuids[i]);
+ }
+ }
+}
+
+} // namespace starrocks::workgroup