Skip to content

Commit 662a918

Browse files
sparknacksparknack
committed
fix: drain index entry futures after errors
Signed-off-by: Shawn Wang <shawn.wang@zilliz.com>
1 parent 2534454 commit 662a918

2 files changed

Lines changed: 149 additions & 44 deletions

File tree

‎internal/core/src/storage/IndexEntryReader.cpp‎

Lines changed: 107 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
#include <cstdint>
2525
#include <cstring>
2626
#include <deque>
27+
#include <exception>
2728
#include <functional>
2829
#include <future>
2930
#include <limits>
@@ -50,6 +51,36 @@ struct ActiveSliceTask {
5051
std::future<void> future;
5152
};
5253

54+
void
55+
RememberFirstError(std::exception_ptr& first_error, std::exception_ptr error) {
56+
if (error && !first_error) {
57+
first_error = std::move(error);
58+
}
59+
}
60+
61+
std::exception_ptr
62+
WaitForAllFutures(std::vector<std::future<void>>& futures) {
63+
std::exception_ptr first_error = nullptr;
64+
for (auto& future : futures) {
65+
if (!future.valid()) {
66+
continue;
67+
}
68+
try {
69+
future.get();
70+
} catch (...) {
71+
RememberFirstError(first_error, std::current_exception());
72+
}
73+
}
74+
return first_error;
75+
}
76+
77+
void
78+
RethrowIfError(const std::exception_ptr& error) {
79+
if (error) {
80+
std::rethrow_exception(error);
81+
}
82+
}
83+
5384
void
5485
ReadOrderedEntryStream(
5586
size_t num_slices,
@@ -435,26 +466,33 @@ IndexEntryReader::ReadPlainEntry(const EntryMeta& meta) {
435466
std::vector<std::future<void>> futures;
436467
size_t remaining = pm.size;
437468
size_t offset = 0;
469+
std::exception_ptr first_error = nullptr;
438470

439-
while (remaining > 0) {
440-
size_t len = std::min(remaining, kRangeSize);
441-
size_t this_offset = offset;
471+
try {
472+
auto input = input_;
473+
size_t entry_offset = pm.offset;
474+
while (remaining > 0) {
475+
size_t len = std::min(remaining, kRangeSize);
476+
size_t this_offset = offset;
442477

443-
futures.push_back(pool.Submit([this, dest, this_offset, len, &pm]() {
444-
size_t n =
445-
input_->ReadAt(dest + this_offset,
446-
MILVUS_V3_MAGIC_SIZE + pm.offset + this_offset,
447-
len);
448-
AssertInfo(n == len, "Failed to read entry data range");
449-
}));
478+
futures.push_back(
479+
pool.Submit([input, dest, this_offset, len, entry_offset]() {
480+
size_t n = input->ReadAt(
481+
dest + this_offset,
482+
MILVUS_V3_MAGIC_SIZE + entry_offset + this_offset,
483+
len);
484+
AssertInfo(n == len, "Failed to read entry data range");
485+
}));
450486

451-
remaining -= len;
452-
offset += len;
487+
remaining -= len;
488+
offset += len;
489+
}
490+
} catch (...) {
491+
first_error = std::current_exception();
453492
}
454493

455-
for (auto& f : futures) {
456-
f.get();
457-
}
494+
RememberFirstError(first_error, WaitForAllFutures(futures));
495+
RethrowIfError(first_error);
458496

459497
// CRC verification: sequential pass over the assembled buffer
460498
VerifyCrc32c(pm.crc32, result.data.data(), pm.size, "");
@@ -473,23 +511,37 @@ IndexEntryReader::ReadEncryptedEntry(const EntryMeta& meta) {
473511

474512
std::vector<std::future<void>> futures;
475513
size_t cur_output_offset = 0;
514+
std::exception_ptr first_error = nullptr;
476515

477-
for (const auto& slice : em.slices) {
478-
size_t this_output_offset = cur_output_offset;
479-
size_t remaining = em.original_size - cur_output_offset;
480-
size_t plain_len = std::min(remaining, slice_size_);
481-
cur_output_offset += plain_len;
516+
try {
517+
auto input = input_;
518+
auto cipher_plugin = cipher_plugin_;
519+
int64_t ez_id = ez_id_;
520+
int64_t collection_id = collection_id_;
521+
auto edek = edek_;
522+
for (const auto& slice : em.slices) {
523+
size_t this_output_offset = cur_output_offset;
524+
size_t remaining = em.original_size - cur_output_offset;
525+
size_t plain_len = std::min(remaining, slice_size_);
526+
cur_output_offset += plain_len;
482527

483-
futures.push_back(
484-
pool.Submit([this, slice, dest, this_output_offset, plain_len]() {
528+
futures.push_back(pool.Submit([input,
529+
cipher_plugin,
530+
ez_id,
531+
collection_id,
532+
edek,
533+
slice,
534+
dest,
535+
this_output_offset,
536+
plain_len]() {
485537
std::vector<uint8_t> cipher(slice.size);
486-
size_t n = input_->ReadAt(cipher.data(),
487-
MILVUS_V3_MAGIC_SIZE + slice.offset,
488-
slice.size);
538+
size_t n = input->ReadAt(cipher.data(),
539+
MILVUS_V3_MAGIC_SIZE + slice.offset,
540+
slice.size);
489541
AssertInfo(n == slice.size, "Failed to read encrypted slice");
490542

491543
auto dec =
492-
cipher_plugin_->GetDecryptor(ez_id_, collection_id_, edek_);
544+
cipher_plugin->GetDecryptor(ez_id, collection_id, edek);
493545
auto plain = dec->Decrypt(cipher.data(), cipher.size());
494546

495547
AssertInfo(plain.size() == plain_len,
@@ -499,11 +551,13 @@ IndexEntryReader::ReadEncryptedEntry(const EntryMeta& meta) {
499551
std::memcpy(
500552
dest + this_output_offset, plain.data(), plain.size());
501553
}));
554+
}
555+
} catch (...) {
556+
first_error = std::current_exception();
502557
}
503558

504-
for (auto& f : futures) {
505-
f.get();
506-
}
559+
RememberFirstError(first_error, WaitForAllFutures(futures));
560+
RethrowIfError(first_error);
507561

508562
// CRC verification over full plaintext buffer
509563
VerifyCrc32c(em.crc32, result.data.data(), em.original_size, "");
@@ -557,6 +611,11 @@ IndexEntryReader::SubmitEntryDownloadTasks(
557611
if (meta.encrypted) {
558612
const auto& em = meta.enc;
559613
size_t output_offset = 0;
614+
auto input = input_;
615+
auto cipher_plugin = cipher_plugin_;
616+
int64_t ez_id = ez_id_;
617+
int64_t collection_id = collection_id_;
618+
auto edek = edek_;
560619

561620
for (size_t i = 0; i < em.slices.size(); i++) {
562621
const auto& slice = em.slices[i];
@@ -565,21 +624,25 @@ IndexEntryReader::SubmitEntryDownloadTasks(
565624
size_t plain_len = std::min(remaining, slice_size_);
566625
output_offset += plain_len;
567626

568-
futures.push_back(pool.Submit([this,
627+
futures.push_back(pool.Submit([input,
628+
cipher_plugin,
629+
ez_id,
630+
collection_id,
631+
edek,
569632
slice,
570633
fd = state.fd,
571634
this_output_offset,
572635
plain_len,
573636
i,
574637
&state]() {
575638
std::vector<uint8_t> cipher(slice.size);
576-
size_t n = input_->ReadAt(cipher.data(),
577-
MILVUS_V3_MAGIC_SIZE + slice.offset,
578-
slice.size);
639+
size_t n = input->ReadAt(cipher.data(),
640+
MILVUS_V3_MAGIC_SIZE + slice.offset,
641+
slice.size);
579642
AssertInfo(n == slice.size, "Failed to read encrypted slice");
580643

581644
auto dec =
582-
cipher_plugin_->GetDecryptor(ez_id_, collection_id_, edek_);
645+
cipher_plugin->GetDecryptor(ez_id, collection_id, edek);
583646
auto plain = dec->Decrypt(cipher.data(), cipher.size());
584647

585648
AssertInfo(plain.size() == plain_len,
@@ -602,22 +665,23 @@ IndexEntryReader::SubmitEntryDownloadTasks(
602665
size_t file_offset = 0;
603666
size_t src_offset = pm.offset;
604667
size_t range_idx = 0;
668+
auto input = input_;
605669

606670
while (remaining > 0) {
607671
size_t len = std::min(remaining, kRangeSize);
608672
size_t this_file_offset = file_offset;
609673
size_t this_src_offset = src_offset;
610674
size_t this_range_idx = range_idx;
611675

612-
futures.push_back(pool.Submit([this,
676+
futures.push_back(pool.Submit([input,
613677
this_src_offset,
614678
len,
615679
fd = state.fd,
616680
this_file_offset,
617681
this_range_idx,
618682
&state]() {
619683
std::vector<uint8_t> buf(len);
620-
size_t n = input_->ReadAt(
684+
size_t n = input->ReadAt(
621685
buf.data(), MILVUS_V3_MAGIC_SIZE + this_src_offset, len);
622686
AssertInfo(n == len, "Failed to read data for file");
623687
auto written = ::pwrite(fd, buf.data(), len, this_file_offset);
@@ -663,18 +727,18 @@ IndexEntryReader::ReadEntryToFile(const std::string& name,
663727
const auto& meta = it->second;
664728

665729
auto state = PrepareEntryDownload(name, local_path, meta);
730+
std::vector<std::future<void>> futures;
666731
try {
667-
std::vector<std::future<void>> futures;
668732
SubmitEntryDownloadTasks(meta, state, futures);
669733

670-
for (auto& f : futures) {
671-
f.get();
672-
}
734+
RethrowIfError(WaitForAllFutures(futures));
673735

674736
FinalizeEntryDownload(state);
675737
} catch (...) {
738+
WaitForAllFutures(futures);
676739
if (state.fd != -1) {
677740
::close(state.fd);
741+
state.fd = -1;
678742
}
679743
throw;
680744
}
@@ -701,6 +765,7 @@ IndexEntryReader::ReadEntriesToFiles(
701765
}
702766
};
703767

768+
std::vector<std::future<void>> all_futures;
704769
try {
705770
for (const auto& [name, path] : name_path_pairs) {
706771
auto it = entry_index_.find(name);
@@ -709,22 +774,20 @@ IndexEntryReader::ReadEntriesToFiles(
709774
}
710775

711776
// Submit ALL tasks for ALL entries at once (avoids thread pool deadlock)
712-
std::vector<std::future<void>> all_futures;
713777
for (size_t i = 0; i < name_path_pairs.size(); i++) {
714778
const auto& meta = entry_index_.at(name_path_pairs[i].first);
715779
SubmitEntryDownloadTasks(meta, states[i], all_futures);
716780
}
717781

718782
// Wait for ALL tasks to complete
719-
for (auto& f : all_futures) {
720-
f.get();
721-
}
783+
RethrowIfError(WaitForAllFutures(all_futures));
722784

723785
// Verify CRCs and close all file descriptors
724786
for (auto& state : states) {
725787
FinalizeEntryDownload(state);
726788
}
727789
} catch (...) {
790+
WaitForAllFutures(all_futures);
728791
close_all_fds();
729792
throw;
730793
}

‎internal/core/src/storage/IndexEntryWriterTest.cpp‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#include <unistd.h>
1515

1616
#include <algorithm>
17+
#include <atomic>
1718
#include <chrono>
1819
#include <cstddef>
1920
#include <cstdint>
@@ -164,6 +165,7 @@ class DelayedFailingInputStream : public milvus::InputStream {
164165
size_t offset;
165166
std::chrono::milliseconds delay;
166167
bool fail;
168+
std::shared_ptr<std::atomic<int>> completion_count = nullptr;
167169
};
168170

169171
DelayedFailingInputStream(std::shared_ptr<milvus::InputStream> base,
@@ -201,6 +203,10 @@ class DelayedFailingInputStream : public milvus::InputStream {
201203
for (const auto& rule : rules_) {
202204
if (rule.offset == offset) {
203205
std::this_thread::sleep_for(rule.delay);
206+
if (rule.completion_count) {
207+
rule.completion_count->fetch_add(1,
208+
std::memory_order_relaxed);
209+
}
204210
if (rule.fail) {
205211
return 0;
206212
}
@@ -731,6 +737,42 @@ TEST_F(IndexEntryWriterV3Test, ReadEntriesToFilesMultiRangeParallel) {
731737
::unlink(file_c.c_str());
732738
}
733739

740+
TEST_F(IndexEntryWriterV3Test, ReadEntriesToFilesDrainsFuturesAfterReadError) {
741+
const std::string file_path = kV3FilePath + "_tofiles_drain_error";
742+
constexpr size_t kRangeSize = 16 * 1024 * 1024;
743+
const size_t entry_size = 2 * kRangeSize + 1024;
744+
auto data = GeneratePattern(entry_size);
745+
746+
{
747+
auto output = CreateOutputStream(file_path);
748+
IndexEntryDirectStreamWriter writer(output);
749+
writer.WriteEntry("entry", data.data(), data.size());
750+
writer.Finish();
751+
}
752+
753+
auto slow_read_done = std::make_shared<std::atomic<int>>(0);
754+
auto base_input = CreateInputStream(file_path);
755+
auto input = std::make_shared<DelayedFailingInputStream>(
756+
base_input,
757+
std::vector<DelayedFailingInputStream::Rule>{
758+
{MILVUS_V3_MAGIC_SIZE, std::chrono::milliseconds(0), true},
759+
{MILVUS_V3_MAGIC_SIZE + kRangeSize,
760+
std::chrono::milliseconds(200),
761+
false,
762+
slow_read_done},
763+
});
764+
int64_t file_size = GetFileSize(file_path);
765+
auto reader = IndexEntryReader::Open(input, file_size);
766+
767+
std::string local_file = GetRootPath() + "/drain_error_output.bin";
768+
std::vector<std::pair<std::string, std::string>> pairs = {
769+
{"entry", local_file}};
770+
EXPECT_THROW(reader->ReadEntriesToFiles(pairs), milvus::SegcoreError);
771+
EXPECT_EQ(slow_read_done->load(std::memory_order_relaxed), 1);
772+
773+
::unlink(local_file.c_str());
774+
}
775+
734776
// =============================================================================
735777
// Edge Case Tests: Directory Table and Meta Entry size boundaries
736778
// =============================================================================

0 commit comments

Comments
 (0)