Skip to content

Commit f3eea47

Browse files
gengjun-gitmergify[bot]
gengjun-git
authored and
mergify[bot]
committed
[BugFix] Cancel pipeline fragments when reaping expired external scan contexts (#76535)
Signed-off-by: gengjun-git <gengjun@starrocks.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> (cherry picked from commit d58c924) # Conflicts: # be/src/orchestration/orchestration_env.cpp # be/src/runtime/external_scan_context_mgr.cpp # be/src/runtime/external_scan_context_mgr.h # be/test/orchestration/external_scan_context_mgr_test.cpp # be/test/orchestration/external_scan_orchestrator_test.cpp
1 parent 21a5952 commit f3eea47

5 files changed

Lines changed: 646 additions & 0 deletions

File tree

‎be/src/orchestration/orchestration_env.cpp‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
// Copyright 2021-present StarRocks, Inc. All rights reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include "orchestration/orchestration_env.h"
16+
17+
#include <unistd.h>
18+
19+
#include <memory>
20+
#include <vector>
21+
22+
#include "common/config_exec_env_fwd.h"
23+
#include "common/logging.h"
24+
#include "common/process_exit.h"
25+
#include "common/system/master_info.h"
26+
#include "compute_env/compute_env.h"
27+
#include "compute_env/profile_report_worker.h"
28+
#include "exec/exec_env.h"
29+
#include "exec/pipeline/pipeline_fragment_reporter.h"
30+
#include "exec/runtime/query_context_manager.h"
31+
#include "orchestration/external_scan_context_mgr.h"
32+
#include "orchestration/external_scan_orchestrator.h"
33+
#include "orchestration/fragment_mgr.h"
34+
#include "orchestration/orchestration_metrics.h"
35+
#include "orchestration/routine_load_task_executor.h"
36+
#include "orchestration/runtime_filter_worker.h"
37+
#include "orchestration/stream_load_orchestrator.h"
38+
39+
namespace starrocks::orchestration {
40+
41+
OrchestrationEnv::OrchestrationEnv() = default;
42+
43+
OrchestrationEnv::~OrchestrationEnv() {
44+
destroy();
45+
}
46+
47+
Status OrchestrationEnv::init(ExecEnv* exec_env, MetricRegistry* metrics, StreamLoadExecutor* stream_load_executor) {
48+
DCHECK(exec_env != nullptr);
49+
DCHECK(stream_load_executor != nullptr);
50+
_exec_env = exec_env;
51+
52+
_fragment_mgr = std::make_unique<FragmentMgr>(exec_env, metrics);
53+
54+
ProfileReportWorkerOptions profile_report_worker_options;
55+
profile_report_worker_options.report_non_pipeline_fragments =
56+
[this](const std::vector<TUniqueId>& non_pipeline_need_report_fragment_ids) {
57+
DCHECK(_fragment_mgr != nullptr);
58+
return _fragment_mgr->report_fragments(non_pipeline_need_report_fragment_ids);
59+
};
60+
profile_report_worker_options.report_pipeline_fragments =
61+
[exec_env](const std::vector<PipeLineReportTaskKey>& pipeline_need_report_query_fragment_ids) {
62+
DCHECK(exec_env->query_context_mgr() != nullptr);
63+
return report_pipeline_fragments(exec_env->query_context_mgr(),
64+
pipeline_need_report_query_fragment_ids);
65+
};
66+
RETURN_IF_ERROR(exec_env->compute_env()->init_profile_report_worker(std::move(profile_report_worker_options)));
67+
68+
_runtime_filter_worker =
69+
std::make_unique<RuntimeFilterWorker>(&exec_env->runtime_services(), &exec_env->rpc_services(),
70+
exec_env->query_pool_mem_tracker(), _fragment_mgr.get());
71+
_runtime_filter_worker_started = true;
72+
exec_env->set_runtime_filter_services(_runtime_filter_worker.get(), _runtime_filter_worker.get());
73+
74+
_metrics = std::make_unique<OrchestrationMetrics>();
75+
_metrics->install(
76+
metrics, [this] { return _runtime_filter_worker == nullptr ? nullptr : _runtime_filter_worker->metrics(); },
77+
[this] { return _runtime_filter_worker == nullptr ? 0 : _runtime_filter_worker->queue_size(); });
78+
79+
_external_scan_context_mgr = std::make_unique<ExternalScanContextMgr>(exec_env, metrics);
80+
_external_scan_orchestrator =
81+
std::make_unique<ExternalScanOrchestrator>(exec_env, _external_scan_context_mgr.get());
82+
83+
_stream_load_orchestrator = std::make_unique<StreamLoadOrchestrator>(exec_env, _fragment_mgr.get());
84+
85+
_routine_load_task_executor =
86+
std::make_unique<RoutineLoadTaskExecutor>(exec_env, _stream_load_orchestrator.get(), stream_load_executor);
87+
RETURN_IF_ERROR(_routine_load_task_executor->init(metrics));
88+
_routine_load_task_executor_started = true;
89+
90+
return Status::OK();
91+
}
92+
93+
void OrchestrationEnv::wait_for_finish() {
94+
if (config::loop_count_wait_fragments_finish < 0) {
95+
LOG(WARNING) << "'config::loop_count_wait_fragments_finish' is set to a negative integer, ignore it.";
96+
return;
97+
}
98+
99+
size_t max_loop_secs = config::loop_count_wait_fragments_finish * 10;
100+
if (max_loop_secs == 0) {
101+
return;
102+
}
103+
104+
size_t running_fragments = _get_running_fragments_count();
105+
size_t loop_secs = 0;
106+
107+
// TODO: decouple the heartbeat with the graceful exit
108+
// only wait for frontend's heartbeat when the node is ever received heartbeats from the frontend
109+
bool need_wait_frontend_hb = config::graceful_exit_wait_for_frontend_heartbeat && get_backend_id().has_value();
110+
111+
while ((running_fragments > 0 || (need_wait_frontend_hb && !is_frontend_aware_of_exit())) &&
112+
loop_secs < max_loop_secs) {
113+
LOG(INFO) << "Frontend is aware of exit: " << is_frontend_aware_of_exit() << ", " << running_fragments
114+
<< " fragment(s) are still running...";
115+
sleep(1);
116+
running_fragments = _get_running_fragments_count();
117+
loop_secs++;
118+
}
119+
}
120+
121+
void OrchestrationEnv::stop() {
122+
if (_exec_env != nullptr && _exec_env->compute_env() != nullptr && _exec_env->profile_report_worker() != nullptr) {
123+
_exec_env->compute_env()->stop_profile_report_worker();
124+
}
125+
if (_runtime_filter_worker != nullptr && _runtime_filter_worker_started) {
126+
_runtime_filter_worker->close();
127+
_runtime_filter_worker_started = false;
128+
}
129+
if (_routine_load_task_executor != nullptr && _routine_load_task_executor_started) {
130+
_routine_load_task_executor->stop();
131+
_routine_load_task_executor_started = false;
132+
}
133+
if (_fragment_mgr != nullptr) {
134+
_fragment_mgr->close();
135+
}
136+
}
137+
138+
void OrchestrationEnv::destroy() {
139+
stop();
140+
_metrics.reset();
141+
_runtime_filter_worker.reset();
142+
_routine_load_task_executor.reset();
143+
_stream_load_orchestrator.reset();
144+
_external_scan_orchestrator.reset();
145+
_external_scan_context_mgr.reset();
146+
if (_exec_env != nullptr && _exec_env->compute_env() != nullptr) {
147+
_exec_env->compute_env()->destroy_profile_report_worker();
148+
_exec_env->set_runtime_filter_services(nullptr, nullptr);
149+
_exec_env = nullptr;
150+
}
151+
_fragment_mgr.reset();
152+
}
153+
154+
size_t OrchestrationEnv::_get_running_fragments_count() const {
155+
const auto non_pipeline_fragments = _fragment_mgr == nullptr ? 0 : _fragment_mgr->running_fragment_count();
156+
const auto pipeline_fragments = (_exec_env == nullptr || _exec_env->query_context_mgr() == nullptr)
157+
? 0
158+
: _exec_env->query_context_mgr()->size();
159+
return non_pipeline_fragments + pipeline_fragments;
160+
}
161+
162+
} // namespace starrocks::orchestration

‎be/src/runtime/external_scan_context_mgr.cpp‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,15 +40,25 @@
4040
#include <vector>
4141

4242
#include "exec/pipeline/query_context.h"
43+
<<<<<<< HEAD:be/src/runtime/external_scan_context_mgr.cpp
4344
#include "runtime/fragment_mgr.h"
4445
#include "runtime/result_queue_mgr.h"
4546
#include "util/starrocks_metrics.h"
4647
#include "util/thread.h"
4748
#include "util/uid_util.h"
49+
=======
50+
#include "exec/runtime/fragment_context_manager.h"
51+
#include "exec/runtime/query_context_manager.h"
52+
#include "runtime/runtime_metrics.h"
53+
>>>>>>> d58c9249a9 ([BugFix] Cancel pipeline fragments when reaping expired external scan contexts (#76535)):be/src/orchestration/external_scan_context_mgr.cpp
4854

4955
namespace starrocks {
5056

57+
<<<<<<< HEAD:be/src/runtime/external_scan_context_mgr.cpp
5158
ExternalScanContextMgr::ExternalScanContextMgr(ExecEnv* exec_env) : _exec_env(exec_env) {
59+
=======
60+
ExternalScanContextMgr::ExternalScanContextMgr(ExecEnv* exec_env, MetricRegistry* metrics) : _exec_env(exec_env) {
61+
>>>>>>> d58c9249a9 ([BugFix] Cancel pipeline fragments when reaping expired external scan contexts (#76535)):be/src/orchestration/external_scan_context_mgr.cpp
5262
// start the reaper thread for gc the expired context
5363
_keep_alive_reaper = std::make_unique<std::thread>(
5464
std::bind<void>(std::mem_fn(&ExternalScanContextMgr::gc_expired_context), this));
@@ -113,6 +123,7 @@ Status ExternalScanContextMgr::clear_scan_context(const std::string& context_id)
113123
}
114124
}
115125
if (context != nullptr) {
126+
<<<<<<< HEAD:be/src/runtime/external_scan_context_mgr.cpp
116127
// cancel pipeline
117128
const auto& fragment_instance_id = context->fragment_instance_id;
118129
if (auto query_ctx = _exec_env->query_context_mgr()->get(context->query_id); query_ctx != nullptr) {
@@ -126,10 +137,33 @@ Status ExternalScanContextMgr::clear_scan_context(const std::string& context_id)
126137
_exec_env->result_queue_mgr()->cancel(fragment_instance_id);
127138
LOG(INFO) << "close scan context: context id [ " << context_id << " ], fragment instance id [ "
128139
<< print_id(fragment_instance_id) << " ]";
140+
=======
141+
LOG(INFO) << "close scan context: context id [ " << context_id << " ], fragment instance id [ "
142+
<< print_id(context->fragment_instance_id) << " ]";
143+
RETURN_IF_ERROR(_cancel_scan_context(context, "close_scanner"));
144+
>>>>>>> d58c9249a9 ([BugFix] Cancel pipeline fragments when reaping expired external scan contexts (#76535)):be/src/orchestration/external_scan_context_mgr.cpp
129145
}
130146
return Status::OK();
131147
}
132148

149+
Status ExternalScanContextMgr::_cancel_scan_context(const std::shared_ptr<ScanContext>& context,
150+
const std::string& reason) {
151+
const auto& fragment_instance_id = context->fragment_instance_id;
152+
// The fragment opened by open_scanner always runs on the pipeline engine, so it must be cancelled
153+
// through the pipeline QueryContext/FragmentContext. Cancelling through the non-pipeline
154+
// FragmentMgr is a silent no-op for pipeline fragments, which leaves an abandoned scanner holding
155+
// all its buffered memory until the query deadline expires.
156+
if (auto query_ctx = _exec_env->query_context_mgr()->get(context->query_id); query_ctx != nullptr) {
157+
if (auto fragment_ctx = query_ctx->fragment_mgr()->get(fragment_instance_id); fragment_ctx != nullptr) {
158+
std::stringstream msg;
159+
msg << "FragmentContext(id=" << print_id(fragment_instance_id) << ") cancelled by " << reason;
160+
pipeline::cancel_fragment_context(fragment_ctx.get(), Status::Cancelled(msg.str()));
161+
}
162+
}
163+
// clear the fragment instance's related result queue
164+
return _exec_env->result_queue_mgr()->cancel(fragment_instance_id);
165+
}
166+
133167
void ExternalScanContextMgr::gc_expired_context() {
134168
#ifndef BE_TEST
135169
while (true) {
@@ -165,11 +199,22 @@ void ExternalScanContextMgr::gc_expired_context() {
165199
}
166200
}
167201
for (const auto& expired_context : expired_contexts) {
202+
<<<<<<< HEAD:be/src/runtime/external_scan_context_mgr.cpp
168203
// must cancel the fragment instance, otherwise return thrift transport TTransportException
169204
_exec_env->fragment_mgr()->cancel(expired_context->fragment_instance_id);
170205
_exec_env->result_queue_mgr()->cancel(expired_context->fragment_instance_id);
206+
=======
207+
WARN_IF_ERROR(_cancel_scan_context(expired_context, "expired scan context gc"),
208+
strings::Substitute("Fail to cancel fragment $0 in result queue mgr",
209+
print_id(expired_context->fragment_instance_id)));
210+
>>>>>>> d58c9249a9 ([BugFix] Cancel pipeline fragments when reaping expired external scan contexts (#76535)):be/src/orchestration/external_scan_context_mgr.cpp
171211
}
172212
}
173213
#endif
174214
}
215+
<<<<<<< HEAD:be/src/runtime/external_scan_context_mgr.cpp
175216
} // namespace starrocks
217+
=======
218+
219+
} // namespace starrocks::orchestration
220+
>>>>>>> d58c9249a9 ([BugFix] Cancel pipeline fragments when reaping expired external scan contexts (#76535)):be/src/orchestration/external_scan_context_mgr.cpp

‎be/src/runtime/external_scan_context_mgr.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,13 +42,21 @@
4242
#include <string>
4343
#include <thread>
4444
#include <utility>
45+
#include <vector>
4546

4647
#include "common/status.h"
4748
#include "gen_cpp/Types_types.h"
4849
#include "runtime/exec_env.h"
4950

5051
namespace starrocks {
5152

53+
<<<<<<< HEAD:be/src/runtime/external_scan_context_mgr.h
54+
=======
55+
class MetricRegistry;
56+
57+
namespace orchestration {
58+
59+
>>>>>>> d58c9249a9 ([BugFix] Cancel pipeline fragments when reaping expired external scan contexts (#76535)):be/src/orchestration/external_scan_context_mgr.h
5260
struct ScanContext {
5361
public:
5462
TUniqueId query_id;
@@ -64,7 +72,11 @@ struct ScanContext {
6472

6573
class ExternalScanContextMgr {
6674
public:
75+
<<<<<<< HEAD:be/src/runtime/external_scan_context_mgr.h
6776
ExternalScanContextMgr(ExecEnv* exec_env);
77+
=======
78+
ExternalScanContextMgr(ExecEnv* exec_env, MetricRegistry* metrics);
79+
>>>>>>> d58c9249a9 ([BugFix] Cancel pipeline fragments when reaping expired external scan contexts (#76535)):be/src/orchestration/external_scan_context_mgr.h
6880

6981
~ExternalScanContextMgr();
7082

@@ -75,6 +87,11 @@ class ExternalScanContextMgr {
7587
Status clear_scan_context(const std::string& context_id);
7688

7789
private:
90+
// Cancel the pipeline fragment of the scan context and clear its result queue. Shared by
91+
// close_scanner (clear_scan_context) and the keep-alive reaper (gc_expired_context); the caller
92+
// must have already removed the context from the active map.
93+
Status _cancel_scan_context(const std::shared_ptr<ScanContext>& context, const std::string& reason);
94+
7895
ExecEnv* _exec_env;
7996
std::map<std::string, std::shared_ptr<ScanContext>> _active_contexts;
8097
void gc_expired_context();

0 commit comments

Comments
 (0)