/root/doris/be/src/runtime/thread_context.cpp
Line | Count | Source (jump to first uncovered line) |
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/thread_context.h" |
19 | | |
20 | | #include "common/signal_handler.h" |
21 | | #include "runtime/query_context.h" |
22 | | #include "runtime/runtime_state.h" |
23 | | |
24 | | namespace doris { |
25 | | class MemTracker; |
26 | | |
27 | 2 | void AttachTask::init(const std::shared_ptr<ResourceContext>& rc) { |
28 | 2 | ThreadLocalHandle::create_thread_local_if_not_exits(); |
29 | 2 | signal::set_signal_task_id(rc->task_controller()->task_id()); |
30 | 2 | thread_context()->attach_task(rc); |
31 | 2 | } |
32 | | |
33 | 2 | AttachTask::AttachTask(const std::shared_ptr<ResourceContext>& rc) { |
34 | 2 | init(rc); |
35 | 2 | } |
36 | | |
37 | 0 | AttachTask::AttachTask(const std::shared_ptr<MemTrackerLimiter>& mem_tracker) { |
38 | 0 | std::shared_ptr<ResourceContext> rc = ResourceContext::create_shared(); |
39 | 0 | rc->memory_context()->set_mem_tracker(mem_tracker); |
40 | 0 | init(rc); |
41 | 0 | } |
42 | | |
43 | 0 | AttachTask::AttachTask(RuntimeState* runtime_state) { |
44 | 0 | signal::set_signal_is_nereids(runtime_state->is_nereids()); |
45 | 0 | init(runtime_state->get_query_ctx()->resource_ctx()); |
46 | 0 | } |
47 | | |
48 | 0 | AttachTask::AttachTask(QueryContext* query_ctx) { |
49 | 0 | init(query_ctx->resource_ctx()); |
50 | 0 | } |
51 | | |
52 | 2 | AttachTask::~AttachTask() { |
53 | 2 | signal::set_signal_task_id(TUniqueId()); |
54 | 2 | thread_context()->detach_task(); |
55 | 2 | ThreadLocalHandle::del_thread_local_if_count_is_zero(); |
56 | 2 | } |
57 | | |
58 | 5 | SwitchResourceContext::SwitchResourceContext(const std::shared_ptr<ResourceContext>& rc) { |
59 | 5 | DCHECK(rc != nullptr); |
60 | 5 | doris::ThreadLocalHandle::create_thread_local_if_not_exits(); |
61 | 5 | if (thread_context()->is_attach_task()) { |
62 | 4 | old_resource_ctx_ = thread_context()->resource_ctx(); |
63 | 4 | if (rc != old_resource_ctx_) { |
64 | 3 | signal::set_signal_task_id(rc->task_controller()->task_id()); |
65 | 3 | thread_context()->resource_ctx_ = rc; |
66 | 3 | thread_context()->thread_mem_tracker_mgr->attach_limiter_tracker( |
67 | 3 | rc->memory_context()->mem_tracker(), rc->workload_group()); |
68 | 3 | } |
69 | 4 | } else { |
70 | 1 | signal::set_signal_task_id(rc->task_controller()->task_id()); |
71 | 1 | thread_context()->attach_task(rc); |
72 | 1 | } |
73 | 5 | } |
74 | | |
75 | 5 | SwitchResourceContext::~SwitchResourceContext() { |
76 | 5 | if (old_resource_ctx_ != thread_context()->resource_ctx()) { |
77 | 4 | if (old_resource_ctx_ != nullptr) { |
78 | 3 | signal::set_signal_task_id(old_resource_ctx_->task_controller()->task_id()); |
79 | 3 | thread_context()->resource_ctx_ = old_resource_ctx_; |
80 | 3 | thread_context()->thread_mem_tracker_mgr->detach_limiter_tracker(); |
81 | 3 | } else { |
82 | 1 | signal::set_signal_task_id(TUniqueId()); |
83 | 1 | thread_context()->detach_task(); |
84 | 1 | } |
85 | 4 | } |
86 | 5 | doris::ThreadLocalHandle::del_thread_local_if_count_is_zero(); |
87 | 5 | } |
88 | | |
89 | | SwitchThreadMemTrackerLimiter::SwitchThreadMemTrackerLimiter( |
90 | 3 | const std::shared_ptr<doris::MemTrackerLimiter>& mem_tracker) { |
91 | 3 | DCHECK(mem_tracker); |
92 | 3 | doris::ThreadLocalHandle::create_thread_local_if_not_exits(); |
93 | 3 | if (mem_tracker != thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker()) { |
94 | 3 | thread_context()->thread_mem_tracker_mgr->attach_limiter_tracker(mem_tracker); |
95 | 3 | is_switched_ = true; |
96 | 3 | } |
97 | 3 | } |
98 | | |
99 | 3 | SwitchThreadMemTrackerLimiter::~SwitchThreadMemTrackerLimiter() { |
100 | 3 | if (is_switched_) { |
101 | 3 | thread_context()->thread_mem_tracker_mgr->detach_limiter_tracker(); |
102 | 3 | } |
103 | 3 | doris::ThreadLocalHandle::del_thread_local_if_count_is_zero(); |
104 | 3 | } |
105 | | |
106 | 0 | AddThreadMemTrackerConsumer::AddThreadMemTrackerConsumer(MemTracker* mem_tracker) { |
107 | 0 | ThreadLocalHandle::create_thread_local_if_not_exits(); |
108 | 0 | if (mem_tracker) { |
109 | 0 | _need_pop = thread_context()->thread_mem_tracker_mgr->push_consumer_tracker(mem_tracker); |
110 | 0 | } |
111 | 0 | } |
112 | | |
113 | | AddThreadMemTrackerConsumer::AddThreadMemTrackerConsumer( |
114 | | const std::shared_ptr<MemTracker>& mem_tracker) |
115 | 3 | : _mem_tracker(mem_tracker) { |
116 | 3 | ThreadLocalHandle::create_thread_local_if_not_exits(); |
117 | 3 | if (_mem_tracker) { |
118 | 3 | _need_pop = |
119 | 3 | thread_context()->thread_mem_tracker_mgr->push_consumer_tracker(_mem_tracker.get()); |
120 | 3 | } |
121 | 3 | } |
122 | | |
123 | 3 | AddThreadMemTrackerConsumer::~AddThreadMemTrackerConsumer() { |
124 | 3 | if (_need_pop) { |
125 | 2 | thread_context()->thread_mem_tracker_mgr->pop_consumer_tracker(); |
126 | 2 | } |
127 | 3 | ThreadLocalHandle::del_thread_local_if_count_is_zero(); |
128 | 3 | } |
129 | | |
130 | | AddThreadMemTrackerConsumerByHook::AddThreadMemTrackerConsumerByHook( |
131 | | const std::shared_ptr<MemTracker>& mem_tracker) |
132 | 0 | : _mem_tracker(mem_tracker) { |
133 | 0 | ThreadLocalHandle::create_thread_local_if_not_exits(); |
134 | 0 | DCHECK(mem_tracker != nullptr); |
135 | 0 | use_mem_hook = true; |
136 | 0 | thread_context()->thread_mem_tracker_mgr->push_consumer_tracker(_mem_tracker.get()); |
137 | 0 | } |
138 | | |
139 | 0 | AddThreadMemTrackerConsumerByHook::~AddThreadMemTrackerConsumerByHook() { |
140 | 0 | thread_context()->thread_mem_tracker_mgr->pop_consumer_tracker(); |
141 | 0 | use_mem_hook = false; |
142 | 0 | ThreadLocalHandle::del_thread_local_if_count_is_zero(); |
143 | 0 | } |
144 | | |
145 | | } // namespace doris |