Coverage Report

Created: 2026-11-03 19:18

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/root/doris/be/src/runtime/query_context.cpp
Line
Count
Source
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
#include "runtime/query_context.h"
19
20
#include <fmt/core.h>
21
#include <gen_cpp/FrontendService_types.h>
22
#include <gen_cpp/RuntimeProfile_types.h>
23
#include <gen_cpp/Types_types.h>
24
#include <glog/logging.h>
25
26
#include <algorithm>
27
#include <exception>
28
#include <memory>
29
#include <mutex>
30
#include <utility>
31
#include <vector>
32
33
#include "cloud/config.h"
34
#include "common/logging.h"
35
#include "common/status.h"
36
#include "exec/operator/rec_cte_scan_operator.h"
37
#include "exec/pipeline/dependency.h"
38
#include "exec/pipeline/pipeline_fragment_context.h"
39
#include "exec/runtime_filter/runtime_filter_definitions.h"
40
#include "exec/spill/spill_file_manager.h"
41
#include "exprs/function/ai/ai_request_executor.h"
42
#include "io/cache/remote_scan_cache_write_limiter.h"
43
#include "io/fs/file_range_read_scheduler.h"
44
#include "runtime/exec_env.h"
45
#include "runtime/fragment_mgr.h"
46
#include "runtime/memory/heap_profiler.h"
47
#include "runtime/runtime_query_statistics_mgr.h"
48
#include "runtime/runtime_state.h"
49
#include "runtime/thread_context.h"
50
#include "runtime/workload_group/workload_group_manager.h"
51
#include "runtime/workload_management/query_task_controller.h"
52
#include "storage/olap_common.h"
53
#include "util/mem_info.h"
54
#include "util/uid_util.h"
55
56
namespace doris {
57
58
class DelayReleaseToken : public Runnable {
59
    ENABLE_FACTORY_CREATOR(DelayReleaseToken);
60
61
public:
62
0
    DelayReleaseToken(std::unique_ptr<ThreadPoolToken>&& token) { token_ = std::move(token); }
63
    ~DelayReleaseToken() override = default;
64
0
    void run() override {}
65
    std::unique_ptr<ThreadPoolToken> token_;
66
};
67
68
0
const std::string toString(QuerySource queryType) {
69
0
    switch (queryType) {
70
0
    case QuerySource::INTERNAL_FRONTEND:
71
0
        return "INTERNAL_FRONTEND";
72
0
    case QuerySource::STREAM_LOAD:
73
0
        return "STREAM_LOAD";
74
0
    case QuerySource::GROUP_COMMIT_LOAD:
75
0
        return "EXTERNAL_QUERY";
76
0
    case QuerySource::ROUTINE_LOAD:
77
0
        return "ROUTINE_LOAD";
78
0
    case QuerySource::EXTERNAL_CONNECTOR:
79
0
        return "EXTERNAL_CONNECTOR";
80
0
    case QuerySource::EXTERNAL_FRONTEND:
81
0
        return "EXTERNAL_FRONTEND";
82
0
    default:
83
0
        return "UNKNOWN";
84
0
    }
85
0
}
86
87
std::shared_ptr<QueryContext> QueryContext::create(TUniqueId query_id, ExecEnv* exec_env,
88
                                                   const TQueryOptions& query_options,
89
                                                   TNetworkAddress coord_addr, bool is_nereids,
90
                                                   TNetworkAddress current_connect_fe,
91
189
                                                   QuerySource query_type) {
92
189
    auto ctx = QueryContext::create_shared(query_id, exec_env, query_options, coord_addr,
93
189
                                           is_nereids, current_connect_fe, query_type);
94
189
    ctx->init_query_task_controller();
95
189
    return ctx;
96
189
}
97
98
QueryContext::QueryContext(TUniqueId query_id, ExecEnv* exec_env,
99
                           const TQueryOptions& query_options, TNetworkAddress coord_addr,
100
                           bool is_nereids, TNetworkAddress current_connect_fe,
101
                           QuerySource query_source)
102
127k
        : _timeout_second(-1),
103
127k
          _query_id(std::move(query_id)),
104
127k
          _exec_env(exec_env),
105
127k
          _is_nereids(is_nereids),
106
127k
          _query_options(query_options),
107
127k
          _query_source(query_source) {
108
127k
    _init_resource_context();
109
127k
    SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(query_mem_tracker());
110
127k
    _query_watcher.start();
111
127k
    _execution_dependency = Dependency::create_unique(-1, -1, "ExecutionDependency", false);
112
127k
    _memory_sufficient_dependency =
113
127k
            Dependency::create_unique(-1, -1, "MemorySufficientDependency", true);
114
115
127k
    _runtime_filter_mgr = std::make_unique<RuntimeFilterMgr>(true);
116
117
127k
    _timeout_second = query_options.execution_timeout;
118
119
127k
    bool initialize_context_holder =
120
127k
            config::enable_file_cache && config::enable_file_cache_query_limit &&
121
127k
            query_options.__isset.enable_file_cache && query_options.enable_file_cache &&
122
127k
            query_options.__isset.file_cache_query_limit_percent &&
123
127k
            query_options.file_cache_query_limit_percent < 100;
124
125
    // Init query context holders for file cache, if enable query limit feature
126
127k
    if (initialize_context_holder) {
127
0
        _query_context_holders = io::FileCacheFactory::instance()->get_query_context_holders(
128
0
                _query_id, query_options.file_cache_query_limit_percent);
129
0
    }
130
131
127k
    const bool initialize_remote_scan_cache_write_limiter =
132
127k
            config::is_cloud_mode() && config::enable_file_cache &&
133
127k
            query_options.__isset.file_cache_query_limit_bytes &&
134
127k
            query_options.file_cache_query_limit_bytes >= 0 &&
135
127k
            query_options.query_type == TQueryType::SELECT;
136
127k
    if (initialize_remote_scan_cache_write_limiter) {
137
3
        _remote_scan_cache_write_limiter = std::make_unique<io::RemoteScanCacheWriteLimiter>(
138
3
                _query_id, query_options.file_cache_query_limit_bytes);
139
3
    }
140
141
127k
    bool is_query_type_valid = query_options.query_type == TQueryType::SELECT ||
142
127k
                               query_options.query_type == TQueryType::LOAD ||
143
127k
                               query_options.query_type == TQueryType::EXTERNAL;
144
127k
    DCHECK_EQ(is_query_type_valid, true);
145
146
127k
    this->coord_addr = coord_addr;
147
    // current_connect_fe is used for report query statistics
148
127k
    this->current_connect_fe = current_connect_fe;
149
    // external query has no current_connect_fe
150
127k
    if (query_options.query_type != TQueryType::EXTERNAL) {
151
343
        bool is_report_fe_addr_valid =
152
343
                !this->current_connect_fe.hostname.empty() && this->current_connect_fe.port != 0;
153
343
        DCHECK_EQ(is_report_fe_addr_valid, true);
154
343
    }
155
127k
    clock_gettime(CLOCK_MONOTONIC, &this->_query_arrival_timestamp);
156
127k
    DorisMetrics::instance()->query_ctx_cnt->increment(1);
157
127k
}
158
159
127k
void QueryContext::_init_query_mem_tracker() {
160
    // If user not set query limit, will use default 1TB memory limit. It is large enough to cover most cases.
161
127k
    constexpr int64_t DEFAULT_QUERY_MEM_LIMIT = 1LL << 60;
162
127k
    bool has_query_mem_limit = _query_options.__isset.mem_limit && (_query_options.mem_limit > 0);
163
127k
    int64_t user_set_mem_limit =
164
127k
            has_query_mem_limit ? _query_options.mem_limit : DEFAULT_QUERY_MEM_LIMIT;
165
127k
    int64_t adjusted_mem_limit = user_set_mem_limit;
166
127k
    if (adjusted_mem_limit > MemInfo::mem_limit()) {
167
39
        VLOG_NOTICE << "Query memory limit "
168
0
                    << PrettyPrinter::print(user_set_mem_limit, TUnit::BYTES)
169
0
                    << " exceeds process memory limit of "
170
0
                    << PrettyPrinter::print(MemInfo::mem_limit(), TUnit::BYTES)
171
0
                    << ". Using process memory limit instead.";
172
39
        adjusted_mem_limit = MemInfo::mem_limit();
173
39
    }
174
    // If the query is a pure load task(streamload, routine load, group commit), then it should not use
175
    // memlimit per query to limit their memory usage.
176
127k
    if (is_pure_load_task()) {
177
127k
        adjusted_mem_limit = MemInfo::mem_limit();
178
127k
    }
179
127k
    std::shared_ptr<MemTrackerLimiter> query_mem_tracker;
180
127k
    if (_query_options.query_type == TQueryType::SELECT) {
181
340
        query_mem_tracker = MemTrackerLimiter::create_shared(
182
340
                MemTrackerLimiter::Type::QUERY, fmt::format("Query#Id={}", print_id(_query_id)),
183
340
                adjusted_mem_limit);
184
127k
    } else if (_query_options.query_type == TQueryType::LOAD) {
185
3
        query_mem_tracker = MemTrackerLimiter::create_shared(
186
3
                MemTrackerLimiter::Type::LOAD, fmt::format("Load#Id={}", print_id(_query_id)),
187
3
                adjusted_mem_limit);
188
127k
    } else if (_query_options.query_type == TQueryType::EXTERNAL) { // spark/flink/etc..
189
127k
        query_mem_tracker = MemTrackerLimiter::create_shared(
190
127k
                MemTrackerLimiter::Type::QUERY, fmt::format("External#Id={}", print_id(_query_id)),
191
127k
                adjusted_mem_limit);
192
127k
    } else {
193
0
        LOG(FATAL) << "__builtin_unreachable";
194
0
        __builtin_unreachable();
195
0
    }
196
127k
    if (_query_options.__isset.is_report_success && _query_options.is_report_success) {
197
0
        query_mem_tracker->enable_print_log_usage();
198
0
    }
199
200
    // If enable reserve memory, not enable check limit, because reserve memory will check it.
201
    // If reserve enabled, even if the reserved memory size is smaller than the actual requested memory,
202
    // and the query memory consumption is larger than the limit, we do not expect the query to fail
203
    // after `check_limit` returns an error, but to run as long as possible,
204
    // and will enter the paused state and try to spill when the query reserves next time.
205
    // If the workload group or process runs out of memory, it will be forced to cancel.
206
127k
    query_mem_tracker->set_enable_check_limit(!(_query_options.__isset.enable_reserve_memory &&
207
127k
                                                _query_options.enable_reserve_memory));
208
127k
    _resource_ctx->memory_context()->set_mem_tracker(query_mem_tracker);
209
127k
    _resource_ctx->memory_context()->set_user_set_mem_limit(user_set_mem_limit);
210
127k
}
211
212
127k
void QueryContext::_init_resource_context() {
213
127k
    _resource_ctx = ResourceContext::create_shared();
214
127k
    _init_query_mem_tracker();
215
127k
}
216
217
127k
void QueryContext::init_query_task_controller() {
218
127k
    _resource_ctx->set_task_controller(QueryTaskController::create(shared_from_this()));
219
127k
    _resource_ctx->task_controller()->set_task_id(_query_id);
220
127k
    _resource_ctx->task_controller()->set_fe_addr(current_connect_fe);
221
127k
    _resource_ctx->task_controller()->set_query_type(_query_options.query_type);
222
#ifndef BE_TEST
223
    _exec_env->runtime_query_statistics_mgr()->register_resource_context(print_id(_query_id),
224
                                                                         _resource_ctx);
225
#endif
226
127k
}
227
228
255
void QueryContext::record_spill_data_dir(SpillDataDir* data_dir) {
229
255
    std::lock_guard lock(_spill_data_dirs_mutex);
230
255
    _spill_data_dirs.emplace(data_dir);
231
255
}
232
233
127k
QueryContext::~QueryContext() {
234
127k
    SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(query_mem_tracker());
235
127k
    {
236
127k
        std::lock_guard lock(_file_range_read_context_mutex);
237
127k
        if (_file_range_read_context != nullptr) {
238
2
            _file_range_read_context->cancel();
239
2
        }
240
127k
    }
241
    // query mem tracker consumption is equal to 0, it means that after QueryContext is created,
242
    // it is found that query already exists in _query_ctx_map, and query mem tracker is not used.
243
    // query mem tracker consumption is not equal to 0 after use, because there is memory consumed
244
    // on query mem tracker, released on other trackers.
245
127k
    std::string mem_tracker_msg;
246
127k
    if (query_mem_tracker()->peak_consumption() != 0) {
247
37
        mem_tracker_msg = fmt::format(
248
37
                "deregister query/load memory tracker, queryId={}, Limit={}, CurrUsed={}, "
249
37
                "PeakUsed={}",
250
37
                print_id(_query_id), PrettyPrinter::print_bytes(query_mem_tracker()->limit()),
251
37
                PrettyPrinter::print_bytes(query_mem_tracker()->consumption()),
252
37
                PrettyPrinter::print_bytes(query_mem_tracker()->peak_consumption()));
253
37
    }
254
127k
    [[maybe_unused]] uint64_t group_id = 0;
255
127k
    if (workload_group()) {
256
46
        group_id = workload_group()->id(); // before remove
257
46
    }
258
259
127k
    _resource_ctx->task_controller()->finish();
260
261
127k
    if (enable_profile()) {
262
0
        _report_query_profile();
263
0
    }
264
265
#ifndef BE_TEST
266
    if (ExecEnv::GetInstance()->pipeline_tracer_context()->enabled()) [[unlikely]] {
267
        try {
268
            ExecEnv::GetInstance()->pipeline_tracer_context()->end_query(_query_id, group_id);
269
        } catch (std::exception& e) {
270
            LOG(WARNING) << "Dump trace log failed bacause " << e.what();
271
        }
272
    }
273
#endif
274
127k
    _runtime_filter_mgr.reset();
275
127k
    _execution_dependency.reset();
276
127k
    _runtime_predicates.clear();
277
127k
    file_scan_range_params_map.clear();
278
127k
    obj_pool.clear();
279
127k
    _merge_controller_handler.reset();
280
281
127k
    if (auto* spill_file_mgr = _exec_env->spill_file_mgr()) {
282
257
        for (auto* data_dir : _spill_data_dirs) {
283
36
            spill_file_mgr->delete_query_spill_directory(print_id(_query_id), data_dir);
284
36
        }
285
257
    }
286
287
127k
    DorisMetrics::instance()->query_ctx_cnt->increment(-1);
288
    // fragment_mgr is nullptr in unittest
289
127k
    if (ExecEnv::GetInstance()->fragment_mgr()) {
290
2
        ExecEnv::GetInstance()->fragment_mgr()->remove_query_context(this->_query_id);
291
2
    }
292
    // the only one msg shows query's end. any other msg should append to it if need.
293
127k
    timespec now;
294
127k
    clock_gettime(CLOCK_MONOTONIC, &now);
295
127k
    int64_t elapsed_ms = (now.tv_sec - _query_arrival_timestamp.tv_sec) * 1000LL +
296
127k
                         (now.tv_nsec - _query_arrival_timestamp.tv_nsec) / 1000000LL;
297
127k
    LOG_INFO("Query {} deconstructed, elapsed_ms: {}, mem_tracker: {}", print_id(this->_query_id),
298
127k
             elapsed_ms, mem_tracker_msg);
299
127k
}
300
301
44
void QueryContext::set_ready_to_execute(Status reason) {
302
44
    set_execution_dependency_ready();
303
44
    _exec_status.update(reason);
304
44
}
305
306
0
void QueryContext::set_ready_to_execute_only() {
307
0
    set_execution_dependency_ready();
308
0
}
309
310
44
void QueryContext::set_execution_dependency_ready() {
311
44
    _execution_dependency->set_ready();
312
44
}
313
314
46
void QueryContext::set_memory_sufficient(bool sufficient) {
315
46
    if (sufficient) {
316
17
        {
317
17
            _memory_sufficient_dependency->set_ready();
318
17
            _resource_ctx->task_controller()->reset_paused_reason();
319
17
        }
320
29
    } else {
321
29
        _memory_sufficient_dependency->block();
322
29
        _resource_ctx->task_controller()->add_paused_count();
323
29
    }
324
46
}
325
326
44
void QueryContext::cancel(Status new_status) {
327
44
    if (!_exec_status.update(new_status)) {
328
0
        return;
329
0
    }
330
44
    {
331
44
        std::lock_guard lock(_file_range_read_context_mutex);
332
44
        if (_file_range_read_context != nullptr) {
333
1
            _file_range_read_context->cancel();
334
1
        }
335
44
    }
336
    // Tasks should be always runnable.
337
44
    _execution_dependency->set_always_ready();
338
44
    _memory_sufficient_dependency->set_always_ready();
339
44
    if ((new_status.is<ErrorCode::MEM_LIMIT_EXCEEDED>() ||
340
44
         new_status.is<ErrorCode::MEM_ALLOC_FAILED>()) &&
341
44
        _query_options.__isset.dump_heap_profile_when_mem_limit_exceeded &&
342
44
        _query_options.dump_heap_profile_when_mem_limit_exceeded) {
343
        // if query is cancelled because of query mem limit exceeded, dump heap profile
344
        // at the time of cancellation can get the most accurate memory usage for problem analysis
345
0
        auto wg = workload_group();
346
0
        auto log_str = fmt::format(
347
0
                "Query {} canceled because of memory limit exceeded, dumping memory "
348
0
                "detail profiles. wg: {}. {}",
349
0
                print_id(_query_id), wg ? wg->debug_string() : "null",
350
0
                doris::ProcessProfile::instance()->memory_profile()->process_memory_detail_str());
351
0
        LOG_LONG_STRING(INFO, log_str);
352
0
        std::string dot = HeapProfiler::instance()->dump_heap_profile_to_dot();
353
0
        if (!dot.empty()) {
354
0
            dot += "\n-------------------------------------------------------\n";
355
0
            dot += "Copy the text after `digraph` in the above output to "
356
0
                   "http://www.webgraphviz.com to generate a dot graph.\n"
357
0
                   "after start heap profiler, if there is no operation, will print `No nodes "
358
0
                   "to "
359
0
                   "print`."
360
0
                   "If there are many errors: `addr2line: Dwarf Error`,"
361
0
                   "or other FAQ, reference doc: "
362
0
                   "https://doris.apache.org/community/developer-guide/debug-tool/#4-qa\n";
363
0
            auto log_str =
364
0
                    fmt::format("Query {}, dump heap profile to dot: {}", print_id(_query_id), dot);
365
0
            LOG_LONG_STRING(INFO, log_str);
366
0
        }
367
0
    }
368
369
44
    set_ready_to_execute(new_status);
370
371
    // Copy the fragment contexts under the map lock, then cancel them after releasing it. Fragment
372
    // cancellation may take task-level locks and must not run while holding the query map lock.
373
44
    std::vector<std::weak_ptr<PipelineFragmentContext>> ctx_to_cancel;
374
44
    {
375
44
        std::lock_guard<std::mutex> lock(_pipeline_map_write_lock);
376
44
        for (auto& entry : _fragment_id_to_pipeline_ctx) {
377
0
            ctx_to_cancel.push_back(entry.second);
378
0
        }
379
44
    }
380
44
    for (auto& f_context : ctx_to_cancel) {
381
0
        if (auto pipeline_ctx = f_context.lock()) {
382
0
            pipeline_ctx->cancel(new_status);
383
0
        }
384
0
    }
385
44
}
386
387
std::shared_ptr<io::FileRangeReadContext> QueryContext::get_or_create_file_range_read_context(
388
3
        io::FileRangeReadScheduler* scheduler) {
389
3
    DORIS_CHECK(scheduler != nullptr);
390
3
    std::lock_guard lock(_file_range_read_context_mutex);
391
3
    if (_file_range_read_context == nullptr) {
392
2
        _file_range_read_scheduler = scheduler;
393
2
        _file_range_read_context = scheduler->create_context();
394
2
        if (is_cancelled()) {
395
1
            _file_range_read_context->cancel();
396
1
        }
397
2
    } else {
398
1
        DORIS_CHECK(_file_range_read_scheduler == scheduler);
399
1
    }
400
3
    return _file_range_read_context;
401
3
}
402
403
0
void QueryContext::set_load_error_url(std::string error_url) {
404
0
    std::lock_guard<std::mutex> lock(_error_url_lock);
405
0
    _load_error_url = error_url;
406
0
}
407
408
0
std::string QueryContext::get_load_error_url() {
409
0
    std::lock_guard<std::mutex> lock(_error_url_lock);
410
0
    return _load_error_url;
411
0
}
412
413
0
void QueryContext::set_first_error_msg(std::string error_msg) {
414
0
    std::lock_guard<std::mutex> lock(_error_url_lock);
415
0
    _first_error_msg = error_msg;
416
0
}
417
418
0
std::string QueryContext::get_first_error_msg() {
419
0
    std::lock_guard<std::mutex> lock(_error_url_lock);
420
0
    return _first_error_msg;
421
0
}
422
423
0
std::string QueryContext::print_all_pipeline_context() {
424
0
    std::vector<std::weak_ptr<PipelineFragmentContext>> ctx_to_print;
425
0
    fmt::memory_buffer debug_string_buffer;
426
0
    size_t i = 0;
427
0
    {
428
0
        fmt::format_to(debug_string_buffer, "{} pipeline fragment contexts in query {}. \n",
429
0
                       _fragment_id_to_pipeline_ctx.size(), print_id(_query_id));
430
431
0
        {
432
0
            std::lock_guard<std::mutex> lock(_pipeline_map_write_lock);
433
0
            for (auto& [f_id, f_context] : _fragment_id_to_pipeline_ctx) {
434
0
                ctx_to_print.push_back(f_context);
435
0
            }
436
0
        }
437
0
        for (auto& f_context : ctx_to_print) {
438
0
            if (auto pipeline_ctx = f_context.lock()) {
439
0
                auto elapsed = pipeline_ctx->elapsed_time() / 1000000000.0;
440
0
                fmt::format_to(debug_string_buffer,
441
0
                               "No.{} (elapse_second={}s, fragment_id={}) : {}\n", i, elapsed,
442
0
                               pipeline_ctx->get_fragment_id(), pipeline_ctx->debug_string());
443
0
                i++;
444
0
            }
445
0
        }
446
0
    }
447
0
    return fmt::to_string(debug_string_buffer);
448
0
}
449
450
28
std::shared_ptr<AIExecutionStatistics> QueryContext::get_ai_statistics(int fragment_id) {
451
28
    std::scoped_lock lock(_ai_statistics_mutex);
452
28
    auto& statistics = _ai_statistics[fragment_id];
453
28
    if (!statistics) {
454
18
        statistics = std::make_shared<AIExecutionStatistics>();
455
18
    }
456
28
    return statistics;
457
28
}
458
459
13
std::shared_ptr<AIExecutionStatistics> QueryContext::find_ai_statistics(int fragment_id) const {
460
13
    std::scoped_lock lock(_ai_statistics_mutex);
461
13
    auto it = _ai_statistics.find(fragment_id);
462
13
    return it == _ai_statistics.end() ? nullptr : it->second;
463
13
}
464
465
void QueryContext::set_pipeline_context(const int fragment_id,
466
0
                                        std::shared_ptr<PipelineFragmentContext> pip_ctx) {
467
0
    std::lock_guard<std::mutex> lock(_pipeline_map_write_lock);
468
    // Use insert_or_assign instead of insert to support overwriting old entries
469
    // when recursive CTE recreates PipelineFragmentContext between rounds.
470
0
    _fragment_id_to_pipeline_ctx.insert_or_assign(fragment_id, pip_ctx);
471
0
}
472
473
30
doris::TaskScheduler* QueryContext::get_pipe_exec_scheduler() {
474
30
    if (!_task_scheduler) {
475
0
        throw Exception(Status::InternalError("task_scheduler is null"));
476
0
    }
477
30
    return _task_scheduler;
478
30
}
479
480
45
Status QueryContext::set_workload_group(WorkloadGroupPtr& wg) {
481
45
    _resource_ctx->set_workload_group(wg);
482
    // Should add query first, the workload group will not be deleted,
483
    // then visit workload group's resource
484
    // see task_group_manager::delete_workload_group_by_ids
485
45
    RETURN_IF_ERROR(workload_group()->add_resource_ctx(_query_id, _resource_ctx));
486
487
45
    workload_group()->get_query_scheduler(&_task_scheduler, &_scan_task_scheduler,
488
45
                                          &_remote_scan_task_scheduler);
489
45
    return Status::OK();
490
45
}
491
492
void QueryContext::add_fragment_profile(
493
        int fragment_id, const std::vector<std::shared_ptr<TRuntimeProfileTree>>& pipeline_profiles,
494
4
        std::shared_ptr<TRuntimeProfileTree> load_channel_profile) {
495
4
    if (pipeline_profiles.empty()) {
496
0
        std::string msg = fmt::format("Add pipeline profile failed, query {}, fragment {}",
497
0
                                      print_id(this->_query_id), fragment_id);
498
0
        LOG_ERROR(msg);
499
0
        DCHECK(false) << msg;
500
0
        return;
501
0
    }
502
503
4
#ifndef NDEBUG
504
4
    for (const auto& p : pipeline_profiles) {
505
4
        DCHECK(p != nullptr) << fmt::format("Add pipeline profile failed, query {}, fragment {}",
506
0
                                            print_id(this->_query_id), fragment_id);
507
4
    }
508
4
#endif
509
510
4
    std::lock_guard<std::mutex> l(_profile_mutex);
511
4
    VLOG_ROW << fmt::format(
512
0
            "Query add fragment profile, query {}, fragment {}, pipeline profile count {} ",
513
0
            print_id(this->_query_id), fragment_id, pipeline_profiles.size());
514
515
4
    _profile_map.insert(std::make_pair(fragment_id, pipeline_profiles));
516
517
4
    if (load_channel_profile != nullptr) {
518
0
        _load_channel_profile_map.insert(std::make_pair(fragment_id, load_channel_profile));
519
0
    }
520
4
}
521
522
0
void QueryContext::_report_query_profile() {
523
0
    std::lock_guard<std::mutex> lg(_profile_mutex);
524
525
0
    for (auto& [fragment_id, fragment_profile] : _profile_map) {
526
        // Recursive CTE stages reuse fragment IDs, but the saved profile is the first stage's
527
        // detached snapshot. Refresh only AI counters from their owner, which accumulates all stages.
528
0
        if (auto statistics = find_ai_statistics(fragment_id)) {
529
0
            const int profile_level =
530
0
                    _query_options.__isset.profile_level ? _query_options.profile_level : 2;
531
0
            statistics->update_profile_snapshot(*fragment_profile.front(), profile_level);
532
0
        }
533
0
        std::shared_ptr<TRuntimeProfileTree> load_channel_profile = nullptr;
534
535
0
        if (_load_channel_profile_map.contains(fragment_id)) {
536
0
            load_channel_profile = _load_channel_profile_map[fragment_id];
537
0
        }
538
539
0
        ExecEnv::GetInstance()->runtime_query_statistics_mgr()->register_fragment_profile(
540
0
                _query_id, this->coord_addr, fragment_id, fragment_profile, load_channel_profile);
541
0
    }
542
543
0
    ExecEnv::GetInstance()->runtime_query_statistics_mgr()->trigger_profile_reporting();
544
0
}
545
546
std::unordered_map<int, std::vector<std::shared_ptr<TRuntimeProfileTree>>>
547
0
QueryContext::_collect_realtime_query_profile() {
548
0
    std::unordered_map<int, std::vector<std::shared_ptr<TRuntimeProfileTree>>> res;
549
0
    std::lock_guard<std::mutex> lock(_pipeline_map_write_lock);
550
0
    for (const auto& [fragment_id, fragment_ctx_wptr] : _fragment_id_to_pipeline_ctx) {
551
0
        if (auto fragment_ctx = fragment_ctx_wptr.lock()) {
552
0
            if (fragment_ctx == nullptr) {
553
0
                std::string msg =
554
0
                        fmt::format("PipelineFragmentContext is nullptr, query {} fragment_id: {}",
555
0
                                    print_id(_query_id), fragment_id);
556
0
                LOG_ERROR(msg);
557
0
                DCHECK(false) << msg;
558
0
                continue;
559
0
            }
560
561
0
            auto profile = fragment_ctx->collect_realtime_profile();
562
563
0
            if (profile.empty()) {
564
0
                std::string err_msg = fmt::format(
565
0
                        "Get nothing when collecting profile, query {}, fragment_id: {}",
566
0
                        print_id(_query_id), fragment_id);
567
0
                LOG_ERROR(err_msg);
568
0
                DCHECK(false) << err_msg;
569
0
                continue;
570
0
            }
571
572
0
            res.insert(std::make_pair(fragment_id, profile));
573
0
        }
574
0
    }
575
576
0
    return res;
577
0
}
578
579
0
TReportExecStatusParams QueryContext::get_realtime_exec_status() {
580
0
    TReportExecStatusParams exec_status;
581
582
0
    auto realtime_query_profile = _collect_realtime_query_profile();
583
0
    std::vector<std::shared_ptr<TRuntimeProfileTree>> load_channel_profiles;
584
585
0
    for (auto load_channel_profile : _load_channel_profile_map) {
586
0
        if (load_channel_profile.second != nullptr) {
587
0
            load_channel_profiles.push_back(load_channel_profile.second);
588
0
        }
589
0
    }
590
591
0
    exec_status = RuntimeQueryStatisticsMgr::create_report_exec_status_params(
592
0
            this->_query_id, std::move(realtime_query_profile), std::move(load_channel_profiles),
593
0
            /*is_done=*/false);
594
595
0
    return exec_status;
596
0
}
597
598
Status QueryContext::send_block_to_cte_scan(
599
        const TUniqueId& instance_id, int node_id,
600
0
        const google::protobuf::RepeatedPtrField<doris::PBlock>& pblocks, bool eos) {
601
0
    std::unique_lock<std::mutex> l(_cte_scan_lock);
602
0
    auto it = _cte_scan.find(std::make_pair(instance_id, node_id));
603
0
    if (it == _cte_scan.end()) {
604
0
        return Status::InternalError("RecCTEScan not found for instance {}, node {}",
605
0
                                     print_id(instance_id), node_id);
606
0
    }
607
0
    for (const auto& pblock : pblocks) {
608
0
        RETURN_IF_ERROR(it->second->add_block(pblock));
609
0
    }
610
0
    if (eos) {
611
0
        it->second->set_ready();
612
0
    }
613
0
    return Status::OK();
614
0
}
615
616
void QueryContext::registe_cte_scan(const TUniqueId& instance_id, int node_id,
617
0
                                    RecCTEScanLocalState* scan) {
618
0
    std::unique_lock<std::mutex> l(_cte_scan_lock);
619
0
    auto key = std::make_pair(instance_id, node_id);
620
0
    DCHECK(!_cte_scan.contains(key)) << "Duplicate registe cte scan for instance "
621
0
                                     << print_id(instance_id) << ", node " << node_id;
622
0
    _cte_scan.emplace(key, scan);
623
0
}
624
625
0
void QueryContext::deregiste_cte_scan(const TUniqueId& instance_id, int node_id) {
626
0
    std::lock_guard<std::mutex> l(_cte_scan_lock);
627
0
    auto key = std::make_pair(instance_id, node_id);
628
0
    DCHECK(_cte_scan.contains(key)) << "Duplicate deregiste cte scan for instance "
629
0
                                    << print_id(instance_id) << ", node " << node_id;
630
0
    _cte_scan.erase(key);
631
0
}
632
633
0
Status QueryContext::reset_global_rf(const google::protobuf::RepeatedField<int32_t>& filter_ids) {
634
0
    if (_merge_controller_handler) {
635
0
        return _merge_controller_handler->reset_global_rf(this, filter_ids);
636
0
    }
637
0
    return Status::OK();
638
0
}
639
640
6
void QueryContext::add_total_task_num(int delta) {
641
6
    if (auto* qtc = dynamic_cast<QueryTaskController*>(_resource_ctx->task_controller())) {
642
6
        qtc->add_total_task_num(delta);
643
6
    }
644
6
}
645
646
421
void QueryContext::inc_finished_task_num() {
647
421
    if (auto* qtc = dynamic_cast<QueryTaskController*>(_resource_ctx->task_controller())) {
648
421
        qtc->inc_finished_task_num();
649
421
    }
650
421
}
651
652
} // namespace doris