Skip to content

Commit fd9c971

Browse files
sollhuisollhui
authored
branch-4.0: [fix](load) replace tablet writer close polling with event wakeup (#64993)
pick #64221
1 parent 89892c3 commit fd9c971

8 files changed

Lines changed: 145 additions & 16 deletions

File tree

‎be/src/vec/sink/load_stream_map_pool.cpp‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@ LoadStreamMap::LoadStreamMap(UniqueId load_id, int64_t src_id, int num_streams,
3232
_num_incremental_streams(0),
3333
_pool(pool),
3434
_tablet_schema_for_index(std::make_shared<IndexToTabletSchema>()),
35-
_enable_unique_mow_for_index(std::make_shared<IndexToEnableMoW>()) {
35+
_enable_unique_mow_for_index(std::make_shared<IndexToEnableMoW>()),
36+
_close_wait_notifier(std::make_shared<CloseWaitNotifier>()) {
3637
DCHECK(num_streams > 0) << "stream num should be greater than 0";
3738
DCHECK(num_use > 0) << "use num should be greater than 0";
3839
}
@@ -46,9 +47,9 @@ std::shared_ptr<LoadStreamStubs> LoadStreamMap::get_or_create(int64_t dst_id, bo
4647
if (incremental) {
4748
_num_incremental_streams.fetch_add(1);
4849
}
49-
streams = std::make_shared<LoadStreamStubs>(_num_streams, _load_id, _src_id,
50-
_tablet_schema_for_index,
51-
_enable_unique_mow_for_index, incremental);
50+
streams = std::make_shared<LoadStreamStubs>(
51+
_num_streams, _load_id, _src_id, _tablet_schema_for_index, _enable_unique_mow_for_index,
52+
incremental, _close_wait_notifier);
5253
_streams_for_node[dst_id] = streams;
5354
return streams;
5455
}
@@ -130,6 +131,14 @@ void LoadStreamMap::close_load(bool incremental) {
130131
}
131132
}
132133

134+
int64_t LoadStreamMap::close_wait_version() const {
135+
return _close_wait_notifier->close_wait_version();
136+
}
137+
138+
void LoadStreamMap::wait_for_close_event(int64_t observed_version, int64_t timeout_ms) {
139+
_close_wait_notifier->wait_for_close_event(observed_version, timeout_ms);
140+
}
141+
133142
LoadStreamMapPool::LoadStreamMapPool() = default;
134143

135144
LoadStreamMapPool::~LoadStreamMapPool() = default;

‎be/src/vec/sink/load_stream_map_pool.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,10 @@ class LoadStreamMap {
9898
// only call this method after release() returns true.
9999
void close_load(bool incremental);
100100

101+
int64_t close_wait_version() const;
102+
103+
void wait_for_close_event(int64_t observed_version, int64_t timeout_ms);
104+
101105
std::unordered_map<int64_t, std::shared_ptr<LoadStreamStubs>> get_streams_for_node() {
102106
decltype(_streams_for_node) snapshot;
103107
{
@@ -118,6 +122,7 @@ class LoadStreamMap {
118122
LoadStreamMapPool* _pool = nullptr;
119123
std::shared_ptr<IndexToTabletSchema> _tablet_schema_for_index;
120124
std::shared_ptr<IndexToEnableMoW> _enable_unique_mow_for_index;
125+
std::shared_ptr<CloseWaitNotifier> _close_wait_notifier;
121126

122127
std::mutex _tablets_to_commit_mutex;
123128
std::unordered_map<int64_t, std::unordered_map<int64_t, PTabletID>> _tablets_to_commit;

‎be/src/vec/sink/load_stream_stub.cpp‎

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,24 @@
3131
namespace doris {
3232
#include "common/compile_check_begin.h"
3333

34+
int64_t CloseWaitNotifier::close_wait_version() const {
35+
return _close_wait_version.load(std::memory_order_acquire);
36+
}
37+
38+
void CloseWaitNotifier::wait_for_close_event(int64_t observed_version, int64_t timeout_ms) {
39+
std::unique_lock<bthread::Mutex> lock(_close_wait_mutex);
40+
if (observed_version != close_wait_version()) {
41+
return;
42+
}
43+
static_cast<void>(_close_wait_cv.wait_for(lock, timeout_ms * 1000));
44+
}
45+
46+
void CloseWaitNotifier::notify_close_wait() {
47+
_close_wait_version.fetch_add(1, std::memory_order_acq_rel);
48+
std::lock_guard<bthread::Mutex> lock(_close_wait_mutex);
49+
_close_wait_cv.notify_all();
50+
}
51+
3452
int LoadStreamReplyHandler::on_received_messages(brpc::StreamId id, butil::IOBuf* const messages[],
3553
size_t size) {
3654
auto stub = _stub.lock();
@@ -119,6 +137,7 @@ void LoadStreamReplyHandler::on_closed(brpc::StreamId id) {
119137
return;
120138
}
121139
stub->_is_closed.store(true);
140+
stub->notify_close_wait();
122141
}
123142

124143
inline std::ostream& operator<<(std::ostream& ostr, const LoadStreamReplyHandler& handler) {
@@ -129,12 +148,16 @@ inline std::ostream& operator<<(std::ostream& ostr, const LoadStreamReplyHandler
129148

130149
LoadStreamStub::LoadStreamStub(PUniqueId load_id, int64_t src_id,
131150
std::shared_ptr<IndexToTabletSchema> schema_map,
132-
std::shared_ptr<IndexToEnableMoW> mow_map, bool incremental)
151+
std::shared_ptr<IndexToEnableMoW> mow_map, bool incremental,
152+
std::shared_ptr<CloseWaitNotifier> close_wait_notifier)
133153
: _load_id(load_id),
134154
_src_id(src_id),
135155
_tablet_schema_for_index(schema_map),
136156
_enable_unique_mow_for_index(mow_map),
137-
_is_incremental(incremental) {};
157+
_is_incremental(incremental),
158+
_close_wait_notifier(std::move(close_wait_notifier)) {
159+
DCHECK(_close_wait_notifier != nullptr);
160+
};
138161

139162
LoadStreamStub::~LoadStreamStub() {
140163
if (_is_open.load() && !_is_closed.load()) {
@@ -365,6 +388,10 @@ Status LoadStreamStub::close_finish_check(RuntimeState* state, bool* is_closed)
365388
return Status::OK();
366389
}
367390

391+
void LoadStreamStub::notify_close_wait() {
392+
_close_wait_notifier->notify_close_wait();
393+
}
394+
368395
void LoadStreamStub::cancel(Status reason) {
369396
LOG(WARNING) << *this << " is cancelled because of " << reason;
370397
if (_is_open.load()) {
@@ -376,6 +403,7 @@ void LoadStreamStub::cancel(Status reason) {
376403
_is_cancelled.store(true);
377404
}
378405
_is_closed.store(true);
406+
notify_close_wait();
379407
}
380408

381409
Status LoadStreamStub::_encode_and_send(PStreamHeader& header, std::span<const Slice> data) {

‎be/src/vec/sink/load_stream_stub.h‎

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,20 @@ using IndexToEnableMoW =
8383
std::allocator<phmap::Pair<const int64_t, bool>>, 4,
8484
std::mutex>;
8585

86+
class CloseWaitNotifier {
87+
public:
88+
int64_t close_wait_version() const;
89+
90+
void wait_for_close_event(int64_t observed_version, int64_t timeout_ms);
91+
92+
void notify_close_wait();
93+
94+
private:
95+
std::atomic<int64_t> _close_wait_version {0};
96+
bthread::Mutex _close_wait_mutex;
97+
bthread::ConditionVariable _close_wait_cv;
98+
};
99+
86100
class LoadStreamReplyHandler : public brpc::StreamInputHandler {
87101
public:
88102
LoadStreamReplyHandler(PUniqueId load_id, int64_t dst_id, std::weak_ptr<LoadStreamStub> stub)
@@ -110,12 +124,17 @@ class LoadStreamStub : public std::enable_shared_from_this<LoadStreamStub> {
110124
// construct new stub
111125
LoadStreamStub(PUniqueId load_id, int64_t src_id,
112126
std::shared_ptr<IndexToTabletSchema> schema_map,
113-
std::shared_ptr<IndexToEnableMoW> mow_map, bool incremental = false);
127+
std::shared_ptr<IndexToEnableMoW> mow_map, bool incremental = false,
128+
std::shared_ptr<CloseWaitNotifier> close_wait_notifier =
129+
std::make_shared<CloseWaitNotifier>());
114130

115131
LoadStreamStub(UniqueId load_id, int64_t src_id,
116132
std::shared_ptr<IndexToTabletSchema> schema_map,
117-
std::shared_ptr<IndexToEnableMoW> mow_map, bool incremental = false)
118-
: LoadStreamStub(load_id.to_proto(), src_id, schema_map, mow_map, incremental) {};
133+
std::shared_ptr<IndexToEnableMoW> mow_map, bool incremental = false,
134+
std::shared_ptr<CloseWaitNotifier> close_wait_notifier =
135+
std::make_shared<CloseWaitNotifier>())
136+
: LoadStreamStub(load_id.to_proto(), src_id, schema_map, mow_map, incremental,
137+
std::move(close_wait_notifier)) {};
119138

120139
// for mock this class in UT
121140
#ifdef BE_TEST
@@ -245,6 +264,8 @@ class LoadStreamStub : public std::enable_shared_from_this<LoadStreamStub> {
245264
tablet_load_infos);
246265

247266
private:
267+
void notify_close_wait();
268+
248269
Status _encode_and_send(PStreamHeader& header, std::span<const Slice> data = {});
249270
Status _send_with_buffer(butil::IOBuf& buf, bool sync = false);
250271
Status _send_with_retry(butil::IOBuf& buf);
@@ -283,6 +304,7 @@ class LoadStreamStub : public std::enable_shared_from_this<LoadStreamStub> {
283304
std::unordered_map<int64_t, Status> _failed_tablets;
284305

285306
bool _is_incremental = false;
307+
std::shared_ptr<CloseWaitNotifier> _close_wait_notifier;
286308

287309
bthread::Mutex _write_mutex;
288310
size_t _bytes_written = 0;
@@ -295,12 +317,14 @@ class LoadStreamStubs {
295317
public:
296318
LoadStreamStubs(size_t num_streams, UniqueId load_id, int64_t src_id,
297319
std::shared_ptr<IndexToTabletSchema> schema_map,
298-
std::shared_ptr<IndexToEnableMoW> mow_map, bool incremental = false)
320+
std::shared_ptr<IndexToEnableMoW> mow_map, bool incremental = false,
321+
std::shared_ptr<CloseWaitNotifier> close_wait_notifier =
322+
std::make_shared<CloseWaitNotifier>())
299323
: _is_incremental(incremental) {
300324
_streams.reserve(num_streams);
301325
for (size_t i = 0; i < num_streams; i++) {
302-
_streams.emplace_back(
303-
new LoadStreamStub(load_id, src_id, schema_map, mow_map, incremental));
326+
_streams.emplace_back(new LoadStreamStub(load_id, src_id, schema_map, mow_map,
327+
incremental, close_wait_notifier));
304328
}
305329
}
306330

‎be/src/vec/sink/writer/vtablet_writer.cpp‎

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,8 @@ bvar::PerSecond<bvar::Adder<int64_t>> g_sink_write_rows_per_second("sink_through
104104
bvar::Adder<int64_t> g_sink_load_back_pressure_version_time_ms(
105105
"load_back_pressure_version_time_ms");
106106

107+
static constexpr int64_t CLOSE_WAIT_EVENT_FALLBACK_MS = 1000;
108+
107109
Status IndexChannel::init(RuntimeState* state, const std::vector<TTabletWithPartition>& tablets,
108110
bool incremental) {
109111
SCOPED_CONSUME_MEM_TRACKER(_index_channel_tracker.get());
@@ -311,6 +313,20 @@ static Status cancel_channel_and_check_intolerable_failure(Status status,
311313
return status;
312314
}
313315

316+
void IndexChannel::wait_for_close_event(int64_t observed_version, int64_t timeout_ms) {
317+
std::unique_lock<bthread::Mutex> lock(_close_wait_mutex);
318+
if (observed_version != close_wait_version()) {
319+
return;
320+
}
321+
static_cast<void>(_close_wait_cv.wait_for(lock, timeout_ms * 1000));
322+
}
323+
324+
void IndexChannel::notify_close_wait() {
325+
_close_wait_version.fetch_add(1, std::memory_order_acq_rel);
326+
std::lock_guard<bthread::Mutex> lock(_close_wait_mutex);
327+
_close_wait_cv.notify_all();
328+
}
329+
314330
Status IndexChannel::close_wait(
315331
RuntimeState* state, WriterStats* writer_stats,
316332
std::unordered_map<int64_t, AddBatchCounter>* node_add_batch_counter_map,
@@ -332,6 +348,7 @@ Status IndexChannel::close_wait(
332348
}
333349
}
334350
while (true) {
351+
int64_t close_wait_version = this->close_wait_version();
335352
RETURN_IF_ERROR(check_each_node_channel_close(
336353
&unfinished_node_channel_ids, node_add_batch_counter_map, writer_stats, status));
337354
bool quorum_success = _quorum_success(unfinished_node_channel_ids, need_finish_tablets);
@@ -342,14 +359,15 @@ Status IndexChannel::close_wait(
342359
<< ", load_id: " << print_id(_parent->_load_id);
343360
break;
344361
}
345-
bthread_usleep(1000 * 10);
362+
wait_for_close_event(close_wait_version, CLOSE_WAIT_EVENT_FALLBACK_MS);
346363
}
347364

348365
// 2. wait for all node channel to complete as much as possible
349366
if (!unfinished_node_channel_ids.empty() && need_wait_after_quorum_success) {
350367
int64_t arrival_quorum_success_time = UnixMillis();
351368
int64_t max_wait_time_ms = _calc_max_wait_time_ms(unfinished_node_channel_ids);
352369
while (true) {
370+
int64_t close_wait_version = this->close_wait_version();
353371
RETURN_IF_ERROR(check_each_node_channel_close(&unfinished_node_channel_ids,
354372
node_add_batch_counter_map, writer_stats,
355373
status));
@@ -373,7 +391,8 @@ Status IndexChannel::close_wait(
373391
<< unfinished_node_channel_host_str.str();
374392
break;
375393
}
376-
bthread_usleep(1000 * 10);
394+
wait_for_close_event(close_wait_version, std::min(CLOSE_WAIT_EVENT_FALLBACK_MS,
395+
max_wait_time_ms - elapsed_ms));
377396
}
378397
}
379398
return status;
@@ -874,6 +893,7 @@ void VNodeChannel::_cancel_with_msg(const std::string& msg) {
874893
}
875894
}
876895
_cancelled = true;
896+
_index_channel->notify_close_wait();
877897
}
878898

879899
void VNodeChannel::_refresh_back_pressure_version_wait_time(
@@ -1146,6 +1166,7 @@ void VNodeChannel::_add_block_success_callback(const PTabletWriterAddBlockResult
11461166
}
11471167
}
11481168
_add_batches_finished = true;
1169+
_index_channel->notify_close_wait();
11491170
}
11501171
} else {
11511172
_cancel_with_msg(fmt::format("{}, add batch req success but status isn't ok, err: {}",
@@ -1196,6 +1217,7 @@ void VNodeChannel::_add_block_failed_callback(const WriteBlockCallbackContext& c
11961217
// if this is last rpc, will must set _add_batches_finished. otherwise, node channel's close_wait
11971218
// will be blocked.
11981219
_add_batches_finished = true;
1220+
_index_channel->notify_close_wait();
11991221
}
12001222
}
12011223

‎be/src/vec/sink/writer/vtablet_writer.h‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
#include <google/protobuf/stubs/callback.h>
3232

3333
// IWYU pragma: no_include <bits/chrono.h>
34+
#include <bthread/condition_variable.h>
3435
#include <bthread/mutex.h>
3536

3637
#include <atomic>
@@ -526,6 +527,14 @@ class IndexChannel {
526527
std::unordered_set<int64_t> unfinished_node_channel_ids,
527528
bool need_wait_after_quorum_success);
528529

530+
int64_t close_wait_version() const {
531+
return _close_wait_version.load(std::memory_order_acquire);
532+
}
533+
534+
void wait_for_close_event(int64_t observed_version, int64_t timeout_ms);
535+
536+
void notify_close_wait();
537+
529538
Status check_each_node_channel_close(
530539
std::unordered_set<int64_t>* unfinished_node_channel_ids,
531540
std::unordered_map<int64_t, AddBatchCounter>* node_add_batch_counter_map,
@@ -616,6 +625,10 @@ class IndexChannel {
616625
std::map<int64_t, std::vector<std::pair<int64_t, int64_t>>> _tablets_filtered_rows;
617626

618627
int64_t _start_time = 0;
628+
629+
std::atomic<int64_t> _close_wait_version {0};
630+
bthread::Mutex _close_wait_mutex;
631+
bthread::ConditionVariable _close_wait_cv;
619632
};
620633
} // namespace vectorized
621634
} // namespace doris

‎be/src/vec/sink/writer/vtablet_writer_v2.cpp‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
#include <gen_cpp/Types_types.h>
2525
#include <gen_cpp/internal_service.pb.h>
2626

27+
#include <algorithm>
2728
#include <cstdint>
2829
#include <mutex>
2930
#include <ranges>
@@ -59,6 +60,8 @@ namespace doris::vectorized {
5960

6061
extern bvar::Adder<int64_t> g_sink_load_back_pressure_version_time_ms;
6162

63+
static constexpr int64_t CLOSE_WAIT_EVENT_FALLBACK_MS = 1000;
64+
6265
VTabletWriterV2::VTabletWriterV2(const TDataSink& t_sink, const VExprContextSPtrs& output_exprs,
6366
std::shared_ptr<pipeline::Dependency> dep,
6467
std::shared_ptr<pipeline::Dependency> fin_dep)
@@ -820,6 +823,7 @@ Status VTabletWriterV2::_close_wait(
820823
}
821824
}
822825
while (true) {
826+
int64_t close_wait_version = _load_stream_map->close_wait_version();
823827
RETURN_IF_ERROR(_check_timeout());
824828
RETURN_IF_ERROR(_check_streams_finish(unfinished_streams, status, streams_for_node));
825829
bool quorum_success = _quorum_success(unfinished_streams, need_finish_tablets);
@@ -829,14 +833,15 @@ Status VTabletWriterV2::_close_wait(
829833
<< ", txn_id: " << _txn_id << ", load_id: " << print_id(_load_id);
830834
break;
831835
}
832-
bthread_usleep(1000 * 10);
836+
_load_stream_map->wait_for_close_event(close_wait_version, CLOSE_WAIT_EVENT_FALLBACK_MS);
833837
}
834838

835839
// 2. then wait for remaining streams as much as possible
836840
if (!unfinished_streams.empty() && need_wait_after_quorum_success) {
837841
int64_t arrival_quorum_success_time = UnixMillis();
838842
int64_t max_wait_time_ms = _calc_max_wait_time_ms(streams_for_node, unfinished_streams);
839843
while (true) {
844+
int64_t close_wait_version = _load_stream_map->close_wait_version();
840845
RETURN_IF_ERROR(_check_timeout());
841846
RETURN_IF_ERROR(_check_streams_finish(unfinished_streams, status, streams_for_node));
842847
if (unfinished_streams.empty()) {
@@ -855,7 +860,9 @@ Status VTabletWriterV2::_close_wait(
855860
<< ", unfinished streams: " << unfinished_streams_str.str();
856861
break;
857862
}
858-
bthread_usleep(1000 * 10);
863+
_load_stream_map->wait_for_close_event(
864+
close_wait_version,
865+
std::min(CLOSE_WAIT_EVENT_FALLBACK_MS, max_wait_time_ms - elapsed_ms));
859866
}
860867
}
861868

0 commit comments

Comments
 (0)