MPP: fix report status error when local RU is not initialized

dbms/src/Flash/Mpp/MPPTaskStatistics.cpp

@@ -95,11 +95,13 @@ void MPPTaskStatistics::collectRuntimeStatistics()
 
 tipb::SelectResponse MPPTaskStatistics::genExecutionSummaryResponse()
 {
+    syncRUInfoToExecutorStatisticsCollector();
     return executor_statistics_collector.genExecutionSummaryResponse();
 }
 
 tipb::TiFlashExecutionInfo MPPTaskStatistics::genTiFlashExecutionInfo()
 {
+    syncRUInfoToExecutorStatisticsCollector();
     return executor_statistics_collector.genTiFlashExecutionInfo();
 }
 
@@ -144,10 +146,15 @@ void MPPTaskStatistics::setMemoryPeak(Int64 memory_peak_)
     memory_peak = memory_peak_;
 }
 
+void MPPTaskStatistics::syncRUInfoToExecutorStatisticsCollector()
+{
+    executor_statistics_collector.setLocalRUConsumption(ru_info);
+}
+
 void MPPTaskStatistics::setRUInfo(const RUConsumption & ru_info_)
 {
     ru_info = ru_info_;
-    executor_statistics_collector.setLocalRUConsumption(ru_info_);
+    syncRUInfoToExecutorStatisticsCollector();
 }
 
 void MPPTaskStatistics::setCompileTimestamp(const Timestamp & start_timestamp, const Timestamp & end_timestamp)

dbms/src/Flash/Mpp/MPPTaskStatistics.h

@@ -62,6 +62,8 @@ class MPPTaskStatistics
     tipb::TiFlashExecutionInfo genTiFlashExecutionInfo();
 
 private:
+    void syncRUInfoToExecutorStatisticsCollector();
+
     void recordInputBytes(DAGContext & dag_context);
 
     const LoggerPtr log;

dbms/src/Flash/Statistics/ExecutorStatisticsCollector.cpp

@@ -153,7 +153,14 @@ tipb::TiFlashExecutionInfo ExecutorStatisticsCollector::genTiFlashExecutionInfo(
 
 void ExecutorStatisticsCollector::setLocalRUConsumption(const RUConsumption & ru_info)
 {
-    local_ru = std::make_optional<resource_manager::Consumption>();
+    if (!local_ru)
+    {
+        local_ru.emplace();
+    }
+    else
+    {
+        local_ru->Clear();
+    }
     local_ru->set_r_r_u(ru_info.cpu_ru + ru_info.read_ru);
     local_ru->set_total_cpu_time_ms(toCPUTimeMillisecond(ru_info.cpu_time_ns));
     local_ru->set_read_bytes(ru_info.read_bytes);

dbms/src/Flash/tests/gtest_execution_summary.cpp

@@ -12,9 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Flash/Mpp/MPPTaskStatistics.h>
+#include <Flash/Mpp/MPPTunnelSet.h>
 #include <Interpreters/Context.h>
 #include <TestUtils/ExecutorTestUtils.h>
 #include <TestUtils/mockExecutor.h>
+#include <kvproto/resource_manager.pb.h>
 
 namespace DB
 {
@@ -150,6 +153,48 @@ try
 }
 CATCH
 
+TEST_F(ExecutionSummaryTestRunner, genMPPTaskExecutionInfoWithoutSetRUInfo)
+try
+{
+    auto request = context.scan("test_db", "test_table").limit(1).exchangeSender(tipb::PassThrough).build(context);
+    request->set_collect_execution_summaries(true);
+
+    mpp::TaskMeta meta;
+    meta.set_gather_id(1);
+    meta.set_task_id(1);
+    meta.set_query_ts(1);
+    meta.set_local_query_id(1);
+    meta.set_server_id(1);
+    meta.set_start_ts(1);
+
+    DAGContext dag_context(*request, meta, false);
+    // ExchangeSenderStatistics reads tunnel_set while initializing the mock MPP DAG.
+    dag_context.tunnel_set = std::make_shared<MPPTunnelSet>("test-host");
+    MPPTaskStatistics task_statistics(MPPTaskId(meta), "test-host");
+    task_statistics.initializeExecutorDAG(&dag_context);
+
+    auto execution_info = task_statistics.genTiFlashExecutionInfo();
+    const auto root_executor_id = dag_context.dag_request.rootExecutorID();
+    const tipb::ExecutorExecutionSummary * root_summary = nullptr;
+    for (const auto & summary : execution_info.execution_summaries())
+    {
+        if (summary.executor_id() == root_executor_id)
+        {
+            root_summary = &summary;
+            break;
+        }
+    }
+
+    ASSERT_NE(root_summary, nullptr);
+    ASSERT_TRUE(root_summary->has_ru_consumption());
+    resource_manager::Consumption ru_consumption;
+    ASSERT_TRUE(ru_consumption.ParseFromString(root_summary->ru_consumption()));
+    ASSERT_DOUBLE_EQ(0.0, ru_consumption.r_r_u());
+    ASSERT_DOUBLE_EQ(0.0, ru_consumption.read_bytes());
+    ASSERT_DOUBLE_EQ(0.0, ru_consumption.total_cpu_time_ms());
+}
+CATCH
+
 TEST_F(ExecutionSummaryTestRunner, treeBased)
 try
 {

dbms/src/Server/Server.cpp

@@ -1051,6 +1051,9 @@ try
           *  table engines could use Context on destroy.
           */
         LOG_INFO(log, "Shutting down storages.");
+        // Stop scheduler first to avoid use-after-free: schedLoop may try to
+        // push MergedTask to already-destroyed reader_pools.
+        DB::DM::SegmentReadTaskScheduler::instance().stop();
         // `SegmentReader` threads may hold a segment and its delta-index for read.
         // `Context::shutdown()` will destroy `DeltaIndexManager`.
         // So, stop threads explicitly before `TiFlashTestEnv::shutdown()`.

dbms/src/Storages/DeltaMerge/ReadThread/SegmentReadTaskScheduler.cpp

@@ -30,6 +30,11 @@ SegmentReadTaskScheduler::SegmentReadTaskScheduler(bool run_sched_thread)
 }
 
 SegmentReadTaskScheduler::~SegmentReadTaskScheduler()
+{
+    stop();
+}
+
+void SegmentReadTaskScheduler::stop()
 {
     setStop();
     if (likely(sched_thread.joinable()))
@@ -38,6 +43,15 @@ SegmentReadTaskScheduler::~SegmentReadTaskScheduler()
     }
 }
 
+void SegmentReadTaskScheduler::pushMergedTask(const MergedTaskPtr & p)
+{
+    // After stop(), the schedLoop is no longer running and merged_task_pool
+    // will never be drained. Discard re-queued tasks to avoid resource leak.
+    if (isStop())
+        return;
+    merged_task_pool.push(p);
+}
+
 void SegmentReadTaskScheduler::add(const SegmentReadTaskPoolPtr & pool)
 {
     // To avoid schedule from always failing to acquire the pending_mtx.
@@ -223,12 +237,12 @@ std::optional<std::pair<GlobalSegmentID, std::vector<UInt64>>> SegmentReadTaskSc
 
 void SegmentReadTaskScheduler::setStop()
 {
-    stop.store(true, std::memory_order_relaxed);
+    stop_flag.store(true, std::memory_order_relaxed);
 }
 
 bool SegmentReadTaskScheduler::isStop() const
 {
-    return stop.load(std::memory_order_relaxed);
+    return stop_flag.load(std::memory_order_relaxed);
 }
 
 std::tuple<UInt64, UInt64, UInt64> SegmentReadTaskScheduler::scheduleOneRound()

dbms/src/Storages/DeltaMerge/ReadThread/SegmentReadTaskScheduler.h

@@ -33,7 +33,7 @@ class SegmentReadTasksPoolTest;
 //
 // - `sched_thread` will scheduling read tasks.
 // - Call path: schedLoop -> schedule -> reapPendingPools -> scheduleOneRound
-// - reapPeningPools will swap the `pending_pools` and add these pools to `read_pools` and `merging_segments`.
+// - reapPendingPools will swap the `pending_pools` and add these pools to `read_pools` and `merging_segments`.
 // - scheduleOneRound will scan `read_pools` and choose segments to read.
 class SegmentReadTaskScheduler
 {
@@ -45,12 +45,17 @@ class SegmentReadTaskScheduler
     }
 
     ~SegmentReadTaskScheduler();
+
+    // Must be called before SegmentReaderPoolManager::stop() during shutdown to avoid
+    // use-after-free: schedLoop may still try to push MergedTask to destroyed reader_pools.
+    void stop();
+
     DISALLOW_COPY_AND_MOVE(SegmentReadTaskScheduler);
 
     // Add `pool` to `pending_pools`.
     void add(const SegmentReadTaskPoolPtr & pool);
 
-    void pushMergedTask(const MergedTaskPtr & p) { merged_task_pool.push(p); }
+    void pushMergedTask(const MergedTaskPtr & p);
 
     void updateConfig(const Settings & settings);
 
@@ -96,7 +101,7 @@ class SegmentReadTaskScheduler
 
     MergedTaskPool merged_task_pool;
 
-    std::atomic<bool> stop{false};
+    std::atomic<bool> stop_flag{false};
     bool enable_data_sharing{true};
     std::thread sched_thread;
 

dbms/src/Storages/DeltaMerge/tests/gtest_segment_read_task.cpp

@@ -649,6 +649,13 @@ try
 
     // stable
     {
+        // skip_check_segment_update prevents write() from scheduling background
+        // flush/merge tasks, avoiding a race where background placeDeltaIndex or
+        // merge delta holds is_updating and causes mergeDeltaAll() to fail silently.
+        FailPointHelper::enableFailPoint(FailPoints::skip_check_segment_update);
+        auto fp_guard
+            = ext::make_scope_guard([]() { FailPointHelper::disableFailPoint(FailPoints::skip_check_segment_update); });
+
         auto block = DMTestEnv::prepareSimpleWriteBlock(0, 4096, false);
         store->write(*db_context, db_context->getSettingsRef(), block);
         store->mergeDeltaAll(*db_context);

dbms/src/TestUtils/gtests_dbms_main.cpp

@@ -78,13 +78,13 @@ int main(int argc, char ** argv)
     DB::DM::SegmentReaderPoolManager::instance().init(4, 1.0);
     DB::DM::SegmentReadTaskScheduler::instance();
 
-    DB::GlobalThreadPool::initialize(/*max_threads*/ 100, /*max_free_threds*/ 10, /*queue_size*/ 1000);
-    DB::S3FileCachePool::initialize(/*max_threads*/ 20, /*max_free_threds*/ 10, /*queue_size*/ 1000);
-    DB::DataStoreS3Pool::initialize(/*max_threads*/ 20, /*max_free_threds*/ 10, /*queue_size*/ 1000);
-    DB::BuildReadTaskForWNPool::initialize(/*max_threads*/ 20, /*max_free_threds*/ 10, /*queue_size*/ 1000);
-    DB::BuildReadTaskForWNTablePool::initialize(/*max_threads*/ 20, /*max_free_threds*/ 10, /*queue_size*/ 1000);
-    DB::BuildReadTaskPool::initialize(/*max_threads*/ 20, /*max_free_threds*/ 10, /*queue_size*/ 1000);
-    DB::RNWritePageCachePool::initialize(/*max_threads*/ 20, /*max_free_threds*/ 10, /*queue_size*/ 1000);
+    DB::GlobalThreadPool::initialize(/*max_threads*/ 100, /*max_free_threads*/ 10, /*queue_size*/ 1000);
+    DB::S3FileCachePool::initialize(/*max_threads*/ 20, /*max_free_threads*/ 10, /*queue_size*/ 1000);
+    DB::DataStoreS3Pool::initialize(/*max_threads*/ 20, /*max_free_threads*/ 10, /*queue_size*/ 1000);
+    DB::BuildReadTaskForWNPool::initialize(/*max_threads*/ 20, /*max_free_threads*/ 10, /*queue_size*/ 1000);
+    DB::BuildReadTaskForWNTablePool::initialize(/*max_threads*/ 20, /*max_free_threads*/ 10, /*queue_size*/ 1000);
+    DB::BuildReadTaskPool::initialize(/*max_threads*/ 20, /*max_free_threads*/ 10, /*queue_size*/ 1000);
+    DB::RNWritePageCachePool::initialize(/*max_threads*/ 20, /*max_free_threads*/ 10, /*queue_size*/ 1000);
     const auto s3_endpoint = Poco::Environment::get("S3_ENDPOINT", "");
     const auto s3_bucket = Poco::Environment::get("S3_BUCKET", "mockbucket");
     const auto s3_root = Poco::Environment::get("S3_ROOT", "tiflash_ut/");
@@ -117,6 +117,9 @@ int main(int argc, char ** argv)
 
     auto ret = RUN_ALL_TESTS();
 
+    // Stop scheduler first to avoid use-after-free: schedLoop may try to
+    // push MergedTask to already-destroyed reader_pools.
+    DB::DM::SegmentReadTaskScheduler::instance().stop();
     // `SegmentReader` threads may hold a segment and its delta-index for read.
     // `TiFlashTestEnv::shutdown()` will destroy `DeltaIndexManager`.
     // Stop threads explicitly before `TiFlashTestEnv::shutdown()`.