be/src/exec/pipeline/task_queue.cpp

Source
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include "exec/pipeline/task_queue.h"

// IWYU pragma: no_include <bits/chrono.h>
#include <chrono> // IWYU pragma: keep
#include <memory>
#include <string>
#include <vector>

#include "common/logging.h"
#include "exec/pipeline/pipeline_task.h"
#include "runtime/workload_group/workload_group.h"

namespace doris {

PipelineTaskSPtr SubTaskQueue::try_take(bool is_steal) {
    if (_queue.empty()) {
        return nullptr;
    }
    auto task = _queue.front();
    _queue.pop();
    return task;
}

std::queue<PipelineTaskSPtr> SubTaskQueue::take_all() {
    std::queue<PipelineTaskSPtr> tasks;
    tasks.swap(_queue);
    return tasks;
}

////////////////////  PriorityTaskQueue ////////////////////

PriorityTaskQueue::PriorityTaskQueue() : _closed(false) {
    double factor = 1;
    for (int i = SUB_QUEUE_LEVEL - 1; i >= 0; i--) {
        _sub_queues[i].set_level_factor(factor);
        factor *= LEVEL_QUEUE_TIME_FACTOR;
    }
}

void PriorityTaskQueue::close() {
    std::vector<std::queue<PipelineTaskSPtr>> tasks_to_release;
    {
        std::unique_lock<std::mutex> lock(_work_size_mutex);
        _closed = true;
        const auto total_task_size = _total_task_size.exchange(0);
        _wait_task.notify_all();
        DorisMetrics::instance()->pipeline_task_queue_size->increment(
                -static_cast<int64_t>(total_task_size));
        for (auto& sub_queue : _sub_queues) {
            tasks_to_release.emplace_back(sub_queue.take_all());
        }
    }
    for (auto& tasks : tasks_to_release) {
        while (!tasks.empty()) {
            tasks.front()->pop_out_runnable_queue();
            tasks.pop();
        }
    }
}

PipelineTaskSPtr PriorityTaskQueue::_try_take_unprotected(bool is_steal) {
    if (_total_task_size == 0 || _closed) {
        return nullptr;
    }

    double min_vruntime = 0;
    int level = -1;
    for (int i = 0; i < SUB_QUEUE_LEVEL; ++i) {
        double cur_queue_vruntime = _sub_queues[i].get_vruntime();
        if (!_sub_queues[i].empty()) {
            if (level == -1 || cur_queue_vruntime < min_vruntime) {
                level = i;
                min_vruntime = cur_queue_vruntime;
            }
        }
    }
    DCHECK(level != -1);
    _queue_level_min_vruntime = uint64_t(min_vruntime);

    auto task = _sub_queues[level].try_take(is_steal);
    if (task) {
        task->update_queue_level(level);
        _total_task_size--;
        DorisMetrics::instance()->pipeline_task_queue_size->increment(-1);
    }
    return task;
}

int PriorityTaskQueue::_compute_level(uint64_t runtime) {
    for (int i = 0; i < SUB_QUEUE_LEVEL - 1; ++i) {
        if (runtime <= _queue_level_limit[i]) {
            return i;
        }
    }
    return SUB_QUEUE_LEVEL - 1;
}

PipelineTaskSPtr PriorityTaskQueue::try_take(bool is_steal) {
    // TODO other efficient lock? e.g. if get lock fail, return null_ptr
    std::unique_lock<std::mutex> lock(_work_size_mutex);
    return _try_take_unprotected(is_steal);
}

PipelineTaskSPtr PriorityTaskQueue::take(uint32_t timeout_ms) {
    std::unique_lock<std::mutex> lock(_work_size_mutex);
    auto task = _try_take_unprotected(false);
    if (task) {
        return task;
    } else {
        if (timeout_ms > 0) {
            _wait_task.wait_for(lock, std::chrono::milliseconds(timeout_ms));
        } else {
            _wait_task.wait(lock);
        }
        return _try_take_unprotected(false);
    }
}

Status PriorityTaskQueue::push(PipelineTaskSPtr task) {
    auto level = _compute_level(task->get_runtime_ns());
    std::unique_lock<std::mutex> lock(_work_size_mutex);
    if (_closed) {
        return Status::InternalError("WorkTaskQueue closed");
    }

    // update empty queue's  runtime, to avoid too high priority
    if (_sub_queues[level].empty() &&
        double(_queue_level_min_vruntime) > _sub_queues[level].get_vruntime()) {
        _sub_queues[level].adjust_runtime(_queue_level_min_vruntime);
    }

    task->put_in_runnable_queue();
    _sub_queues[level].push_back(task);
    _total_task_size++;
    DorisMetrics::instance()->pipeline_task_queue_size->increment(1);
    _wait_task.notify_one();
    return Status::OK();
}

MultiCoreTaskQueue::~MultiCoreTaskQueue() = default;

MultiCoreTaskQueue::MultiCoreTaskQueue(int core_size)
        : _prio_task_queues(core_size), _closed(false), _core_size(core_size) {}

void MultiCoreTaskQueue::close() {
    if (_closed) {
        return;
    }
    _closed = true;
    // close all priority task queue
    std::ranges::for_each(_prio_task_queues,
                          [](auto& prio_task_queue) { prio_task_queue.close(); });
}

PipelineTaskSPtr MultiCoreTaskQueue::take(int core_id) {
    PipelineTaskSPtr task = nullptr;
    while (!_closed) {
        DCHECK(_prio_task_queues.size() > core_id)
                << " list size: " << _prio_task_queues.size() << " core_id: " << core_id
                << " _core_size: " << _core_size << " _next_core: " << _next_core.load();
        task = _prio_task_queues[core_id].try_take(false);
        if (task) {
            break;
        }
        task = _steal_take(core_id);
        if (task) {
            break;
        }
        task = _prio_task_queues[core_id].take(WAIT_CORE_TASK_TIMEOUT_MS /* timeout_ms */);
        if (task) {
            break;
        }
    }
    if (task) {
        task->pop_out_runnable_queue();
    }
    return task;
}

PipelineTaskSPtr MultiCoreTaskQueue::_steal_take(int core_id) {
    DCHECK(core_id < _core_size);
    int next_id = core_id;
    for (int i = 1; i < _core_size; ++i) {
        ++next_id;
        if (next_id == _core_size) {
            next_id = 0;
        }
        DCHECK(next_id < _core_size);
        auto task = _prio_task_queues[next_id].try_take(true);
        if (task) {
            return task;
        }
    }
    return nullptr;
}

Status MultiCoreTaskQueue::push_back(PipelineTaskSPtr task) {
    int thread_id = task->get_thread_id(_core_size);
    if (thread_id < 0) {
        thread_id = _next_core.fetch_add(1) % _core_size;
    }
    return push_back(task, thread_id);
}

Status MultiCoreTaskQueue::push_back(PipelineTaskSPtr task, int core_id) {
    DCHECK(core_id < _core_size);
    return _prio_task_queues[core_id].push(task);
}

void MultiCoreTaskQueue::update_statistics(PipelineTask* task, int64_t time_spent) {
    // if the task not execute but exception early close, core_id == -1
    // should not do update_statistics
    if (auto core_id = task->get_thread_id(_core_size); core_id >= 0) {
        task->inc_runtime_ns(time_spent);
        _prio_task_queues[core_id].inc_sub_queue_runtime(task->get_queue_level(), time_spent);
    }
}

} // namespace doris

Coverage Report

Created: 2026-07-04 04:47

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 "exec/pipeline/task_queue.h"
19
20		// IWYU pragma: no_include <bits/chrono.h>
21		#include <chrono> // IWYU pragma: keep
22		#include <memory>
23		#include <string>
24		#include <vector>
25
26		#include "common/logging.h"
27		#include "exec/pipeline/pipeline_task.h"
28		#include "runtime/workload_group/workload_group.h"
29
30		namespace doris {
31
32	12	PipelineTaskSPtr SubTaskQueue::try_take(bool is_steal) {
33	12	if (_queue.empty()) {
34	0	return nullptr;
35	0	}
36	12	auto task = _queue.front();
37	12	_queue.pop();
38	12	return task;
39	12	}
40
41	36	std::queue<PipelineTaskSPtr> SubTaskQueue::take_all() {
42	36	std::queue<PipelineTaskSPtr> tasks;
43	36	tasks.swap(_queue);
44	36	return tasks;
45	36	}
46
47		//////////////////// PriorityTaskQueue ////////////////////
48
49	45	PriorityTaskQueue::PriorityTaskQueue() : _closed(false) {
50	45	double factor = 1;
51	315	for (int i = SUB_QUEUE_LEVEL - 1; i >= 0; i--) {
52	270	_sub_queues[i].set_level_factor(factor);
53	270	factor *= LEVEL_QUEUE_TIME_FACTOR;
54	270	}
55	45	}
56
57	6	void PriorityTaskQueue::close() {
58	6	std::vector<std::queue<PipelineTaskSPtr>> tasks_to_release;
59	6	{
60	6	std::unique_lock<std::mutex> lock(_work_size_mutex);
61	6	_closed = true;
62	6	const auto total_task_size = _total_task_size.exchange(0);
63	6	_wait_task.notify_all();
64	6	DorisMetrics::instance()->pipeline_task_queue_size->increment(
65	6	-static_cast<int64_t>(total_task_size));
66	36	for (auto& sub_queue : _sub_queues) {
67	36	tasks_to_release.emplace_back(sub_queue.take_all());
68	36	}
69	6	}
70	36	for (auto& tasks : tasks_to_release) {
71	36	while (!tasks.empty()) {
72	0	tasks.front()->pop_out_runnable_queue();
73	0	tasks.pop();
74	0	}
75	36	}
76	6	}
77
78	21	PipelineTaskSPtr PriorityTaskQueue::_try_take_unprotected(bool is_steal) {
79	21	if (_total_task_size == 0 \|\| _closed) {
80	9	return nullptr;
81	9	}
82
83	12	double min_vruntime = 0;
84	12	int level = -1;
85	84	for (int i = 0; i < SUB_QUEUE_LEVEL; ++i) {
86	72	double cur_queue_vruntime = _sub_queues[i].get_vruntime();
87	72	if (!_sub_queues[i].empty()) {
88	12	if (level == -1 \|\| cur_queue_vruntime < min_vruntime) {
89	12	level = i;
90	12	min_vruntime = cur_queue_vruntime;
91	12	}
92	12	}
93	72	}
94	12	DCHECK(level != -1);
95	12	_queue_level_min_vruntime = uint64_t(min_vruntime);
96
97	12	auto task = _sub_queues[level].try_take(is_steal);
98	12	if (task) {
99	12	task->update_queue_level(level);
100	12	_total_task_size--;
101	12	DorisMetrics::instance()->pipeline_task_queue_size->increment(-1);
102	12	}
103	12	return task;
104	21	}
105
106	33	int PriorityTaskQueue::_compute_level(uint64_t runtime) {
107	33	for (int i = 0; i < SUB_QUEUE_LEVEL - 1; ++i) {
108	33	if (runtime <= _queue_level_limit[i]) {
109	33	return i;
110	33	}
111	33	}
112	0	return SUB_QUEUE_LEVEL - 1;
113	33	}
114
115	15	PipelineTaskSPtr PriorityTaskQueue::try_take(bool is_steal) {
116		// TODO other efficient lock? e.g. if get lock fail, return null_ptr
117	15	std::unique_lock<std::mutex> lock(_work_size_mutex);
118	15	return _try_take_unprotected(is_steal);
119	15	}
120
121	3	PipelineTaskSPtr PriorityTaskQueue::take(uint32_t timeout_ms) {
122	3	std::unique_lock<std::mutex> lock(_work_size_mutex);
123	3	auto task = _try_take_unprotected(false);
124	3	if (task) {
125	0	return task;
126	3	} else {
127	3	if (timeout_ms > 0) {
128	3	_wait_task.wait_for(lock, std::chrono::milliseconds(timeout_ms));
129	3	} else {
130	0	_wait_task.wait(lock);
131	0	}
132	3	return _try_take_unprotected(false);
133	3	}
134	3	}
135
136	33	Status PriorityTaskQueue::push(PipelineTaskSPtr task) {
137	33	auto level = _compute_level(task->get_runtime_ns());
138	33	std::unique_lock<std::mutex> lock(_work_size_mutex);
139	33	if (_closed) {
140	0	return Status::InternalError("WorkTaskQueue closed");
141	0	}
142
143		// update empty queue's runtime, to avoid too high priority
144	33	if (_sub_queues[level].empty() &&
145	33	double(_queue_level_min_vruntime) > _sub_queues[level].get_vruntime()) {
146	0	_sub_queues[level].adjust_runtime(_queue_level_min_vruntime);
147	0	}
148
149	33	task->put_in_runnable_queue();
150	33	_sub_queues[level].push_back(task);
151	33	_total_task_size++;
152	33	DorisMetrics::instance()->pipeline_task_queue_size->increment(1);
153	33	_wait_task.notify_one();
154	33	return Status::OK();
155	33	}
156
157	87	MultiCoreTaskQueue::~MultiCoreTaskQueue() = default;
158
159		MultiCoreTaskQueue::MultiCoreTaskQueue(int core_size)
160	87	: _prio_task_queues(core_size), _closed(false), _core_size(core_size) {}
161
162	48	void MultiCoreTaskQueue::close() {
163	48	if (_closed) {
164	0	return;
165	0	}
166	48	_closed = true;
167		// close all priority task queue
168	48	std::ranges::for_each(_prio_task_queues,
169	48	[](auto& prio_task_queue) { prio_task_queue.close(); });
170	48	}
171
172	6	PipelineTaskSPtr MultiCoreTaskQueue::take(int core_id) {
173	6	PipelineTaskSPtr task = nullptr;
174	8	while (!_closed) {
175	3	DCHECK(_prio_task_queues.size() > core_id)
176	0	<< " list size: " << _prio_task_queues.size() << " core_id: " << core_id
177	0	<< " _core_size: " << _core_size << " _next_core: " << _next_core.load();
178	3	task = _prio_task_queues[core_id].try_take(false);
179	3	if (task) {
180	1	break;
181	1	}
182	2	task = _steal_take(core_id);
183	2	if (task) {
184	0	break;
185	0	}
186	2	task = _prio_task_queues[core_id].take(WAIT_CORE_TASK_TIMEOUT_MS /* timeout_ms */);
187	2	if (task) {
188	0	break;
189	0	}
190	2	}
191	6	if (task) {
192	1	task->pop_out_runnable_queue();
193	1	}
194	6	return task;
195	6	}
196
197	3	PipelineTaskSPtr MultiCoreTaskQueue::_steal_take(int core_id) {
198	3	DCHECK(core_id < _core_size);
199	3	int next_id = core_id;
200	3	for (int i = 1; i < _core_size; ++i) {
201	0	++next_id;
202	0	if (next_id == _core_size) {
203	0	next_id = 0;
204	0	}
205	0	DCHECK(next_id < _core_size);
206	0	auto task = _prio_task_queues[next_id].try_take(true);
207	0	if (task) {
208	0	return task;
209	0	}
210	0	}
211	3	return nullptr;
212	3	}
213
214	33	Status MultiCoreTaskQueue::push_back(PipelineTaskSPtr task) {
215	33	int thread_id = task->get_thread_id(_core_size);
216	33	if (thread_id < 0) {
217	33	thread_id = _next_core.fetch_add(1) % _core_size;
218	33	}
219	33	return push_back(task, thread_id);
220	33	}
221
222	33	Status MultiCoreTaskQueue::push_back(PipelineTaskSPtr task, int core_id) {
223	33	DCHECK(core_id < _core_size);
224	33	return _prio_task_queues[core_id].push(task);
225	33	}
226
227	1	void MultiCoreTaskQueue::update_statistics(PipelineTask* task, int64_t time_spent) {
228		// if the task not execute but exception early close, core_id == -1
229		// should not do update_statistics
230	1	if (auto core_id = task->get_thread_id(_core_size); core_id >= 0) {
231	1	task->inc_runtime_ns(time_spent);
232	1	_prio_task_queues[core_id].inc_sub_queue_runtime(task->get_queue_level(), time_spent);
233	1	}
234	1	}
235
236		} // namespace doris