Coverage Report

Created: 2025-03-10 17:42

/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