Coverage Report

Created: 2026-03-14 20:54

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
be/src/runtime/result_queue_mgr.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/result_queue_mgr.h"
19
20
#include <gen_cpp/Types_types.h>
21
22
#include <utility>
23
24
#include "common/config.h"
25
#include "common/metrics/doris_metrics.h"
26
#include "common/metrics/metrics.h"
27
#include "common/status.h"
28
#include "runtime/record_batch_queue.h"
29
#include "util/hash_util.hpp"
30
31
namespace doris {
32
33
DEFINE_GAUGE_METRIC_PROTOTYPE_2ARG(result_block_queue_count, MetricUnit::NOUNIT);
34
35
10
ResultQueueMgr::ResultQueueMgr() {
36
    // Each BlockingQueue has a limited size (default 20, by config::max_memory_sink_batch_count),
37
    // it's not needed to count the actual size of all BlockingQueue.
38
10
    REGISTER_HOOK_METRIC(result_block_queue_count, [this]() {
39
        // std::lock_guard<std::mutex> l(_lock);
40
10
        return _fragment_queue_map.size();
41
10
    });
42
10
}
43
44
10
ResultQueueMgr::~ResultQueueMgr() {
45
10
    DEREGISTER_HOOK_METRIC(result_block_queue_count);
46
10
}
47
48
Status ResultQueueMgr::fetch_result(const TUniqueId& fragment_instance_id,
49
2
                                    std::shared_ptr<arrow::RecordBatch>* result, bool* eos) {
50
2
    BlockQueueSharedPtr queue;
51
2
    {
52
2
        std::lock_guard<std::mutex> l(_lock);
53
2
        auto iter = _fragment_queue_map.find(fragment_instance_id);
54
2
        if (_fragment_queue_map.end() != iter) {
55
2
            queue = iter->second;
56
2
        } else {
57
0
            return Status::InternalError("fragment_instance_id does not exists");
58
0
        }
59
2
    }
60
    // check queue status before get result
61
2
    RETURN_IF_ERROR(queue->status());
62
2
    bool success = queue->blocking_get(result);
63
2
    if (success) {
64
        // sentinel nullptr indicates scan end
65
2
        if (*result == nullptr) {
66
1
            *eos = true;
67
            // put sentinel for consistency, avoid repeated invoking fetch result when have no rowbatch
68
1
            queue->blocking_put(nullptr);
69
1
        } else {
70
1
            *eos = false;
71
1
        }
72
2
    } else {
73
0
        *eos = true;
74
0
    }
75
2
    return Status::OK();
76
2
}
77
78
void ResultQueueMgr::create_queue(const TUniqueId& fragment_instance_id,
79
7
                                  BlockQueueSharedPtr* queue) {
80
7
    std::lock_guard<std::mutex> l(_lock);
81
7
    auto iter = _fragment_queue_map.find(fragment_instance_id);
82
7
    if (iter != _fragment_queue_map.end()) {
83
1
        *queue = iter->second;
84
6
    } else {
85
        // max_elements will not take effect, because when queue size reaches max_memory_sink_batch_count,
86
        // MemoryScratchSink will block queue dependency, in this way, one queue have 20 * 1024 rows at most.
87
        // use MemoryScratchSink queue dependency instead of BlockingQueue to achieve blocking.
88
6
        BlockQueueSharedPtr tmp(new RecordBatchQueue(config::max_memory_sink_batch_count * 2));
89
6
        _fragment_queue_map.insert(std::make_pair(fragment_instance_id, tmp));
90
6
        *queue = tmp;
91
6
    }
92
7
}
93
94
3
Status ResultQueueMgr::cancel(const TUniqueId& fragment_instance_id) {
95
3
    std::lock_guard<std::mutex> l(_lock);
96
3
    auto iter = _fragment_queue_map.find(fragment_instance_id);
97
3
    if (iter != _fragment_queue_map.end()) {
98
        // first remove RecordBatch from queue
99
        // avoid MemoryScratchSink block on send or close operation
100
2
        iter->second->shutdown();
101
        // remove this queue from map
102
2
        _fragment_queue_map.erase(fragment_instance_id);
103
2
    }
104
3
    return Status::OK();
105
3
}
106
107
void ResultQueueMgr::update_queue_status(const TUniqueId& fragment_instance_id,
108
0
                                         const Status& status) {
109
0
    if (status.ok()) {
110
0
        return;
111
0
    }
112
0
    std::lock_guard<std::mutex> l(_lock);
113
0
    auto iter = _fragment_queue_map.find(fragment_instance_id);
114
0
    if (iter != _fragment_queue_map.end()) {
115
0
        iter->second->update_status(status);
116
0
    }
117
0
}
118
119
} // namespace doris