Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[BugFix] fix load spill writing data size report #55814

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

[BugFix] fix load spill writing data size report #55814

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
90fb1d7
[BugFix] fix load spill writing data size report
luohaha Feb 11, 2025
74501ad
refactor
luohaha Feb 11, 2025
a525ccc
fix ut
luohaha Feb 13, 2025
36fcbc8
update
luohaha Feb 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions be/src/storage/delta_writer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -572,7 +572,8 @@ Status DeltaWriter::flush_memtable_async(bool eos) {
if ((_mem_table != nullptr && _mem_table->get_result_chunk() != nullptr) || eos) {
auto replicate_token = _replicate_token.get();
return _flush_token->submit(
std::move(_mem_table), eos, [replicate_token, this](std::unique_ptr<SegmentPB> seg, bool eos) {
std::move(_mem_table), eos,
[replicate_token, this](SegmentPBPtr seg, bool eos, int64_t flush_data_size) {
if (seg) {
_tablet->add_in_writing_data_size(_opt.txn_id, seg->data_size());
}
Expand All @@ -597,7 +598,7 @@ Status DeltaWriter::flush_memtable_async(bool eos) {
} else {
if (_mem_table != nullptr && _mem_table->get_result_chunk() != nullptr) {
return _flush_token->submit(
std::move(_mem_table), eos, [this](std::unique_ptr<SegmentPB> seg, bool eos) {
std::move(_mem_table), eos, [this](SegmentPBPtr seg, bool eos, int64_t flush_data_size) {
if (seg) {
_tablet->add_in_writing_data_size(_opt.txn_id, seg->data_size());
}
Expand All @@ -616,7 +617,7 @@ Status DeltaWriter::flush_memtable_async(bool eos) {
}
} else if (_replica_state == Peer) {
if (_mem_table != nullptr && _mem_table->get_result_chunk() != nullptr) {
return _flush_token->submit(std::move(_mem_table), eos, [this](std::unique_ptr<SegmentPB> seg, bool eos) {
return _flush_token->submit(std::move(_mem_table), eos, [this](SegmentPBPtr seg, bool eos, int64_t f) {
if (seg) {
_tablet->add_in_writing_data_size(_opt.txn_id, seg->data_size());
}
Expand Down
40 changes: 23 additions & 17 deletions be/src/storage/lake/delta_writer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,13 +55,15 @@ class TabletWriterSink : public MemTableSink {

DISALLOW_COPY_AND_MOVE(TabletWriterSink);

Status flush_chunk(const Chunk& chunk, starrocks::SegmentPB* segment = nullptr, bool eos = false) override {
Status flush_chunk(const Chunk& chunk, starrocks::SegmentPB* segment = nullptr, bool eos = false,
int64_t* flush_data_size = nullptr) override {
RETURN_IF_ERROR(_writer->write(chunk, segment));
return _writer->flush(segment);
}

Status flush_chunk_with_deletes(const Chunk& upserts, const Column& deletes,
starrocks::SegmentPB* segment = nullptr, bool eos = false) override {
starrocks::SegmentPB* segment = nullptr, bool eos = false,
int64_t* flush_data_size = nullptr) override {
RETURN_IF_ERROR(_writer->flush_del_file(deletes));
RETURN_IF_ERROR(_writer->write(upserts, segment));
return _writer->flush(segment);
Expand Down Expand Up @@ -318,21 +320,25 @@ inline Status DeltaWriterImpl::flush_async() {
if (_miss_auto_increment_column && _mem_table->get_result_chunk() != nullptr) {
RETURN_IF_ERROR(fill_auto_increment_id(*_mem_table->get_result_chunk()));
}
st = _flush_token->submit(std::move(_mem_table), _eos, [this](std::unique_ptr<SegmentPB> seg, bool eos) {
if (_immutable_tablet_size > 0 && !_is_immutable.load(std::memory_order_relaxed)) {
if (seg) {
_tablet_manager->add_in_writing_data_size(_tablet_id, seg->data_size());
}
if (_tablet_manager->in_writing_data_size(_tablet_id) > _immutable_tablet_size) {
_is_immutable.store(true, std::memory_order_relaxed);
}
VLOG(2) << "flush memtable, tablet=" << _tablet_id << ", txn=" << _txn_id
<< " _immutable_tablet_size=" << _immutable_tablet_size
<< ", segment_size=" << (seg ? seg->data_size() : 0)
<< ", in_writing_data_size=" << _tablet_manager->in_writing_data_size(_tablet_id)
<< ", is_immutable=" << _is_immutable.load(std::memory_order_relaxed);
}
});
st = _flush_token->submit(
std::move(_mem_table), _eos, [this](SegmentPBPtr seg, bool eos, int64_t flush_data_size) {
if (_immutable_tablet_size > 0 && !_is_immutable.load(std::memory_order_relaxed)) {
if (seg) {
_tablet_manager->add_in_writing_data_size(_tablet_id, seg->data_size());
} else if (flush_data_size > 0) {
// When enable load spill, seg is nullptr, so we need to use flush_data_size
_tablet_manager->add_in_writing_data_size(_tablet_id, flush_data_size);
}
if (_tablet_manager->in_writing_data_size(_tablet_id) > _immutable_tablet_size) {
_is_immutable.store(true, std::memory_order_relaxed);
}
VLOG(2) << "flush memtable, tablet=" << _tablet_id << ", txn=" << _txn_id
<< " _immutable_tablet_size=" << _immutable_tablet_size
<< ", flush data size=" << (seg ? seg->data_size() : flush_data_size)
<< ", in_writing_data_size=" << _tablet_manager->in_writing_data_size(_tablet_id)
<< ", is_immutable=" << _is_immutable.load(std::memory_order_relaxed);
}
});
_mem_table.reset(nullptr);
_last_write_ts = 0;
}
luohaha marked this conversation as resolved.
Show resolved Hide resolved
Expand Down
12 changes: 9 additions & 3 deletions be/src/storage/lake/spill_mem_table_sink.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ Status LoadSpillOutputDataStream::append(RuntimeState* state, const std::vector<
// preallocate block
RETURN_IF_ERROR(_preallocate(total_size));
// append data
_append_bytes += total_size;
return _block->append(data);
}

Expand Down Expand Up @@ -123,7 +124,8 @@ Status SpillMemTableSink::_do_spill(const Chunk& chunk, const spill::SpillOutput
return Status::OK();
}

Status SpillMemTableSink::flush_chunk(const Chunk& chunk, starrocks::SegmentPB* segment, bool eos) {
Status SpillMemTableSink::flush_chunk(const Chunk& chunk, starrocks::SegmentPB* segment, bool eos,
int64_t* flush_data_size) {
if (eos && _block_manager->block_container()->empty()) {
// If there is only one flush, flush it to segment directly
RETURN_IF_ERROR(_writer->write(chunk, segment));
Expand All @@ -137,19 +139,23 @@ Status SpillMemTableSink::flush_chunk(const Chunk& chunk, starrocks::SegmentPB*
RETURN_IF_ERROR(_do_spill(chunk, output));
// 3. flush
RETURN_IF_ERROR(output->flush());
// record append bytes to segment pb
if (flush_data_size != nullptr) {
*flush_data_size = output->append_bytes();
}
return Status::OK();
}

Status SpillMemTableSink::flush_chunk_with_deletes(const Chunk& upserts, const Column& deletes,
starrocks::SegmentPB* segment, bool eos) {
starrocks::SegmentPB* segment, bool eos, int64_t* flush_data_size) {
if (eos && _block_manager->block_container()->empty()) {
// If there is only one flush, flush it to segment directly
RETURN_IF_ERROR(_writer->flush_del_file(deletes));
RETURN_IF_ERROR(_writer->write(upserts, segment));
return _writer->flush(segment);
}
// 1. flush upsert
RETURN_IF_ERROR(flush_chunk(upserts, segment, eos));
RETURN_IF_ERROR(flush_chunk(upserts, segment, eos, flush_data_size));
// 2. flush deletes
RETURN_IF_ERROR(_writer->flush_del_file(deletes));
return Status::OK();
Expand Down
9 changes: 7 additions & 2 deletions be/src/storage/lake/spill_mem_table_sink.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@ class LoadSpillOutputDataStream : public spill::SpillOutputDataStream {

bool is_remote() const override;

int64_t append_bytes() const { return _append_bytes; }

private:
Status _preallocate(size_t block_size);

Expand All @@ -49,17 +51,20 @@ class LoadSpillOutputDataStream : public spill::SpillOutputDataStream {
private:
LoadSpillBlockManager* _block_manager = nullptr;
spill::BlockPtr _block;
int64_t _append_bytes = 0;
};

class SpillMemTableSink : public MemTableSink {
public:
SpillMemTableSink(LoadSpillBlockManager* block_manager, TabletWriter* writer, RuntimeProfile* profile);
~SpillMemTableSink() override = default;

Status flush_chunk(const Chunk& chunk, starrocks::SegmentPB* segment = nullptr, bool eos = false) override;
Status flush_chunk(const Chunk& chunk, starrocks::SegmentPB* segment = nullptr, bool eos = false,
int64_t* flush_data_size = nullptr) override;

Status flush_chunk_with_deletes(const Chunk& upserts, const Column& deletes,
starrocks::SegmentPB* segment = nullptr, bool eos = false) override;
starrocks::SegmentPB* segment = nullptr, bool eos = false,
int64_t* flush_data_size = nullptr) override;

Status merge_blocks_to_segments();

Expand Down
6 changes: 3 additions & 3 deletions be/src/storage/memtable.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -326,7 +326,7 @@ Status MemTable::finalize() {
return Status::OK();
}

Status MemTable::flush(SegmentPB* seg_info, bool eos) {
Status MemTable::flush(SegmentPB* seg_info, bool eos, int64_t* flush_data_size) {
if (UNLIKELY(_result_chunk == nullptr)) {
return Status::OK();
}
Expand All @@ -339,9 +339,9 @@ Status MemTable::flush(SegmentPB* seg_info, bool eos) {
{
SCOPED_RAW_TIMER(&duration_ns);
if (_deletes) {
RETURN_IF_ERROR(_sink->flush_chunk_with_deletes(*_result_chunk, *_deletes, seg_info, eos));
RETURN_IF_ERROR(_sink->flush_chunk_with_deletes(*_result_chunk, *_deletes, seg_info, eos, flush_data_size));
} else {
RETURN_IF_ERROR(_sink->flush_chunk(*_result_chunk, seg_info, eos));
RETURN_IF_ERROR(_sink->flush_chunk(*_result_chunk, seg_info, eos, flush_data_size));
}
}
auto io_stat = scope.current_scoped_tls_io();
Expand Down
2 changes: 1 addition & 1 deletion be/src/storage/memtable.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ class MemTable {
// return true suggests caller should flush this memory table
StatusOr<bool> insert(const Chunk& chunk, const uint32_t* indexes, uint32_t from, uint32_t size);

Status flush(SegmentPB* seg_info = nullptr, bool eos = false);
Status flush(SegmentPB* seg_info = nullptr, bool eos = false, int64_t* flush_data_size = nullptr);

Status finalize();

Expand Down
15 changes: 8 additions & 7 deletions be/src/storage/memtable_flush_executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ namespace starrocks {
class MemtableFlushTask final : public Runnable {
public:
MemtableFlushTask(FlushToken* flush_token, std::unique_ptr<MemTable> memtable, bool eos,
std::function<void(std::unique_ptr<SegmentPB>, bool)> cb)
std::function<void(std::unique_ptr<SegmentPB>, bool, int64_t)> cb)
: _flush_token(flush_token),
_memtable(std::move(memtable)),
_eos(eos),
Expand All @@ -57,12 +57,13 @@ class MemtableFlushTask final : public Runnable {
_flush_token->_stats.queueing_memtable_num--;
_flush_token->_stats.pending_time_ns += MonotonicNanos() - _create_time_ns;
std::unique_ptr<SegmentPB> segment = nullptr;
int64_t flush_data_size = 0;
if (_memtable) {
SCOPED_THREAD_LOCAL_MEM_SETTER(_memtable->mem_tracker(), false);
segment = std::make_unique<SegmentPB>();

_flush_token->_stats.cur_flush_count++;
_flush_token->_flush_memtable(_memtable.get(), segment.get(), _eos);
_flush_token->_flush_memtable(_memtable.get(), segment.get(), _eos, &flush_data_size);
_flush_token->_stats.cur_flush_count--;
_memtable.reset();

Expand All @@ -78,15 +79,15 @@ class MemtableFlushTask final : public Runnable {
}

if (_cb) {
_cb(std::move(segment), _eos);
_cb(std::move(segment), _eos, flush_data_size);
}
}

private:
FlushToken* _flush_token;
std::unique_ptr<MemTable> _memtable;
bool _eos;
std::function<void(std::unique_ptr<SegmentPB>, bool)> _cb;
std::function<void(std::unique_ptr<SegmentPB>, bool, int64_t)> _cb;
int64_t _create_time_ns;
};

Expand All @@ -97,7 +98,7 @@ std::ostream& operator<<(std::ostream& os, const FlushStatistic& stat) {
}

Status FlushToken::submit(std::unique_ptr<MemTable> memtable, bool eos,
std::function<void(std::unique_ptr<SegmentPB>, bool)> cb) {
std::function<void(std::unique_ptr<SegmentPB>, bool, int64_t)> cb) {
RETURN_IF_ERROR(status());
if (memtable == nullptr && !eos) {
return Status::InternalError(fmt::format("memtable=null eos=false"));
Expand Down Expand Up @@ -130,13 +131,13 @@ Status FlushToken::wait() {
return _status;
}

void FlushToken::_flush_memtable(MemTable* memtable, SegmentPB* segment, bool eos) {
void FlushToken::_flush_memtable(MemTable* memtable, SegmentPB* segment, bool eos, int64_t* flush_data_size) {
// If previous flush has failed, return directly
if (!status().ok()) {
return;
}

set_status(memtable->flush(segment, eos));
set_status(memtable->flush(segment, eos, flush_data_size));
_stats.flush_count++;
_stats.memtable_stats += memtable->get_stat();
}
Expand Down
5 changes: 3 additions & 2 deletions be/src/storage/memtable_flush_executor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ struct FlushStatistic {
};

std::ostream& operator<<(std::ostream& os, const FlushStatistic& stat);
using SegmentPBPtr = std::unique_ptr<SegmentPB>;

// A thin wrapper of ThreadPoolToken to submit task.
// For a tablet, there may be multiple memtables, which will be flushed to disk
Expand All @@ -76,7 +77,7 @@ class FlushToken {
: _flush_token(std::move(flush_pool_token)), _status() {}

Status submit(std::unique_ptr<MemTable> mem_table, bool eos = false,
std::function<void(std::unique_ptr<SegmentPB>, bool)> cb = nullptr);
std::function<void(SegmentPBPtr, bool, int64_t)> cb = nullptr);

// error has happpens, so we cancel this token
// And remove all tasks in the queue.
Expand Down Expand Up @@ -109,7 +110,7 @@ class FlushToken {
private:
friend class MemtableFlushTask;

void _flush_memtable(MemTable* memtable, SegmentPB* segment, bool eos);
void _flush_memtable(MemTable* memtable, SegmentPB* segment, bool eos, int64_t* flush_data_size);

std::unique_ptr<ThreadPoolToken> _flush_token;

Expand Down
5 changes: 3 additions & 2 deletions be/src/storage/memtable_rowset_writer_sink.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,12 +28,13 @@ class MemTableRowsetWriterSink : public MemTableSink {

DISALLOW_COPY(MemTableRowsetWriterSink);

Status flush_chunk(const Chunk& chunk, SegmentPB* seg_info = nullptr, bool eos = false) override {
Status flush_chunk(const Chunk& chunk, SegmentPB* seg_info = nullptr, bool eos = false,
int64_t* flush_data_size = nullptr) override {
return _rowset_writer->flush_chunk(chunk, seg_info);
}

Status flush_chunk_with_deletes(const Chunk& upserts, const Column& deletes, SegmentPB* seg_info = nullptr,
bool eos = false) override {
bool eos = false, int64_t* flush_data_size = nullptr) override {
return _rowset_writer->flush_chunk_with_deletes(upserts, deletes, seg_info);
}

Expand Down
5 changes: 3 additions & 2 deletions be/src/storage/memtable_sink.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,10 @@ class MemTableSink {
public:
virtual ~MemTableSink() = default;

virtual Status flush_chunk(const Chunk& chunk, starrocks::SegmentPB* seg_info = nullptr, bool eos = false) = 0;
virtual Status flush_chunk(const Chunk& chunk, starrocks::SegmentPB* seg_info = nullptr, bool eos = false,
int64_t* flush_data_size = nullptr) = 0;
virtual Status flush_chunk_with_deletes(const Chunk& upserts, const Column& deletes, SegmentPB* seg_info = nullptr,
bool eos = false) = 0;
bool eos = false, int64_t* flush_data_size = nullptr) = 0;
};

} // namespace starrocks
5 changes: 2 additions & 3 deletions be/test/storage/lake/async_delta_writer_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -574,9 +574,7 @@ TEST_F(LakeAsyncDeltaWriterTest, test_block_merger) {
CountDownLatch latch(10);
// flush multi times and generate spill blocks
int64_t old_val = config::write_buffer_size;
bool old_val2 = config::enable_load_spill;
config::write_buffer_size = 1;
config::enable_load_spill = true;
ASSIGN_OR_ABORT(auto delta_writer, AsyncDeltaWriterBuilder()
.set_tablet_manager(_tablet_mgr.get())
.set_tablet_id(tablet_id)
Expand All @@ -585,6 +583,7 @@ TEST_F(LakeAsyncDeltaWriterTest, test_block_merger) {
.set_mem_tracker(_mem_tracker.get())
.set_schema_id(_tablet_schema->id())
.set_profile(&_dummy_runtime_profile)
.set_immutable_tablet_size(10000000)
.build());
ASSERT_OK(delta_writer->open());
for (int i = 0; i < 10; i++) {
Expand All @@ -596,14 +595,14 @@ TEST_F(LakeAsyncDeltaWriterTest, test_block_merger) {
}
latch.wait();
config::write_buffer_size = old_val;
config::enable_load_spill = old_val2;
// finish
CountDownLatch latch2(1);
delta_writer->finish([&](StatusOr<TxnLogPtr> res) {
ASSERT_TRUE(res.ok()) << res.ok();
latch2.count_down();
});
latch2.wait();
ASSERT_TRUE(_tablet_mgr->in_writing_data_size(tablet_id) > 0);
}

} // namespace starrocks::lake
6 changes: 3 additions & 3 deletions be/test/storage/memtable_flush_executor_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -327,7 +327,7 @@ TEST_F(MemTableFlushExecutorTest, testMemtableFlushWithSeg) {
bool ret_eos = true;
ASSERT_TRUE(flush_token
->submit(std::move(mem_table), false,
[&](std::unique_ptr<SegmentPB> seg, bool eos) {
[&](std::unique_ptr<SegmentPB> seg, bool eos, int64_t flush_data_size) {
ret_num_rows = seg->num_rows();
ret_eos = eos;
})
Expand Down Expand Up @@ -359,7 +359,7 @@ TEST_F(MemTableFlushExecutorTest, testMemtableFlushWithNullSeg) {
std::unique_ptr<SegmentPB> ret_seg = make_unique<SegmentPB>();
ASSERT_TRUE(flush_token
->submit(std::move(mem_table), true,
[&](std::unique_ptr<SegmentPB> seg, bool eos) {
[&](std::unique_ptr<SegmentPB> seg, bool eos, int64_t flush_data_size) {
ret_seg = std::move(seg);
ret_eos = eos;
})
Expand Down Expand Up @@ -394,7 +394,7 @@ TEST_F(MemTableFlushExecutorTest, testMemtableFlushStatusNotOk) {
flush_token->set_status(Status::NotSupported("Not Suppoted"));
ASSERT_FALSE(flush_token->status().ok());

flush_token->_flush_memtable(nullptr, nullptr, false);
flush_token->_flush_memtable(nullptr, nullptr, false, nullptr);

ASSERT_TRUE(MemTableFlushExecutor::calc_max_threads_for_lake_table(data_dirs) > 0);
}
Expand Down
Loading