Skip to content

Commit d8fd08b

Browse files
zhjwpkuzhjwpku
committed
fix: bound manifest rewrite reads and use shared update ownership
1 parent 8422d6f commit d8fd08b

3 files changed

Lines changed: 288 additions & 51 deletions

File tree

‎src/iceberg/test/rewrite_manifests_test.cc‎

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
#include <algorithm>
2323
#include <format>
2424
#include <memory>
25+
#include <mutex>
2526
#include <optional>
2627
#include <stdexcept>
2728
#include <string>
@@ -36,6 +37,7 @@
3637
#include "iceberg/constants.h"
3738
#include "iceberg/expression/expressions.h"
3839
#include "iceberg/file_io.h"
40+
#include "iceberg/file_reader.h"
3941
#include "iceberg/partition_spec.h"
4042
#include "iceberg/schema.h"
4143
#include "iceberg/table_metadata.h"
@@ -56,6 +58,104 @@ namespace iceberg {
5658

5759
namespace {
5860

61+
struct ManifestReadState {
62+
std::unordered_set<std::string> source_paths;
63+
std::mutex mutex;
64+
int open_readers = 0;
65+
int max_open_readers = 0;
66+
int opened_readers = 0;
67+
int64_t unread_entries = 0;
68+
int64_t max_unread_entries = 0;
69+
int64_t processed_entries = 0;
70+
bool fail_after_first_batch = false;
71+
};
72+
73+
class TrackingManifestReader : public Reader {
74+
public:
75+
TrackingManifestReader(std::unique_ptr<Reader> delegate,
76+
std::shared_ptr<ManifestReadState> state)
77+
: delegate_(std::move(delegate)), state_(std::move(state)) {}
78+
79+
~TrackingManifestReader() override { Release(); }
80+
81+
Status Open(const ReaderOptions& options) override {
82+
auto read_options = options;
83+
const bool track = state_->source_paths.contains(options.path);
84+
if (track) {
85+
read_options.properties.Set(ReaderProperties::kBatchSize, int64_t{2});
86+
}
87+
ICEBERG_RETURN_UNEXPECTED(delegate_->Open(read_options));
88+
if (track) {
89+
std::lock_guard lock(state_->mutex);
90+
tracked_ = true;
91+
++state_->opened_readers;
92+
state_->max_open_readers =
93+
std::max(state_->max_open_readers, ++state_->open_readers);
94+
}
95+
return {};
96+
}
97+
98+
Status Close() override {
99+
Release();
100+
return delegate_->Close();
101+
}
102+
103+
Result<std::optional<ArrowArray>> Next() override {
104+
if (tracked_ && state_->fail_after_first_batch && read_batch_) {
105+
return IOError("Injected streaming manifest read failure");
106+
}
107+
ICEBERG_ASSIGN_OR_RAISE(auto batch, delegate_->Next());
108+
if (tracked_ && batch.has_value()) {
109+
read_batch_ = true;
110+
std::lock_guard lock(state_->mutex);
111+
state_->unread_entries += batch->length;
112+
state_->max_unread_entries =
113+
std::max(state_->max_unread_entries, state_->unread_entries);
114+
}
115+
return batch;
116+
}
117+
118+
Result<ArrowSchema> Schema() override { return delegate_->Schema(); }
119+
120+
Result<std::unordered_map<std::string, std::string>> Metadata() override {
121+
return delegate_->Metadata();
122+
}
123+
124+
private:
125+
void Release() {
126+
if (tracked_) {
127+
std::lock_guard lock(state_->mutex);
128+
--state_->open_readers;
129+
tracked_ = false;
130+
}
131+
}
132+
133+
std::unique_ptr<Reader> delegate_;
134+
std::shared_ptr<ManifestReadState> state_;
135+
bool tracked_ = false;
136+
bool read_batch_ = false;
137+
};
138+
139+
class ScopedManifestReadTracking {
140+
public:
141+
explicit ScopedManifestReadTracking(std::shared_ptr<ManifestReadState> state)
142+
: previous_(ReaderFactoryRegistry::GetFactory(FileFormatType::kAvro)) {
143+
ReaderFactoryRegistry::GetFactory(FileFormatType::kAvro) =
144+
[factory = previous_,
145+
state = std::move(state)]() -> Result<std::unique_ptr<Reader>> {
146+
ICEBERG_ASSIGN_OR_RAISE(auto reader, factory());
147+
return std::make_unique<TrackingManifestReader>(std::move(reader), state);
148+
};
149+
}
150+
151+
~ScopedManifestReadTracking() {
152+
ReaderFactoryRegistry::GetFactory(FileFormatType::kAvro) = std::move(previous_);
153+
}
154+
155+
private:
156+
ReaderFactory previous_;
157+
};
158+
59159
struct ManifestWriteFailureState {
60160
int manifest_output_count = 0;
61161
int fail_on_manifest_output = 0;
@@ -646,6 +746,123 @@ TEST_P(RewriteManifestsTest, RewriteManifestsAppendedDirectly) {
646746
{ManifestStatus::kExisting}, {append_snapshot_id});
647747
}
648748

749+
TEST_P(RewriteManifestsTest, FactoryResultCanCommitDirectly) {
750+
ASSERT_THAT(AppendFiles({file_a_, file_b_}), IsOk());
751+
ICEBERG_UNWRAP_OR_FAIL(auto original, table_->current_snapshot());
752+
ICEBERG_UNWRAP_OR_FAIL(auto ctx,
753+
TransactionContext::Make(table_, TransactionKind::kUpdate));
754+
ICEBERG_UNWRAP_OR_FAIL(auto rewrite,
755+
RewriteManifests::Make(table_->name().name, std::move(ctx)));
756+
rewrite->ClusterBy([](const DataFile&) { return ""; });
757+
ASSERT_THAT(rewrite->Commit(), IsOk());
758+
759+
ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests());
760+
ASSERT_EQ(manifests.size(), 1U);
761+
ExpectManifestEntriesWithSnapshotIds(
762+
manifests[0], {file_a_->file_path, file_b_->file_path},
763+
{ManifestStatus::kExisting, ManifestStatus::kExisting},
764+
{original->snapshot_id, original->snapshot_id});
765+
}
766+
767+
TEST_P(RewriteManifestsTest, StreamsEntriesWithBoundedReaders) {
768+
ICEBERG_UNWRAP_OR_FAIL(auto append, table_->NewFastAppend());
769+
std::vector<std::string> expected_paths;
770+
for (int i = 0; i < 40; ++i) {
771+
std::vector<std::shared_ptr<DataFile>> files;
772+
for (int j = 0; j < 3; ++j) {
773+
auto file = MakeDataFile(std::format("stream-{}-{}", i, j), 1, 10, i);
774+
files.push_back(file);
775+
expected_paths.push_back(file->file_path);
776+
}
777+
ICEBERG_UNWRAP_OR_FAIL(auto manifest,
778+
WriteAddedManifest(std::format("stream-{}", i), files));
779+
append->AppendManifest(manifest);
780+
}
781+
ASSERT_THAT(append->Commit(), IsOk());
782+
783+
for (bool use_executor : {false, true}) {
784+
SCOPED_TRACE(use_executor);
785+
ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests());
786+
auto state = std::make_shared<ManifestReadState>();
787+
for (const auto& manifest : before) {
788+
state->source_paths.insert(manifest.manifest_path);
789+
}
790+
{
791+
test::ThreadExecutor executor;
792+
ScopedManifestReadTracking tracking(state);
793+
ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests());
794+
rewrite->ClusterBy([state](const DataFile& file) {
795+
std::lock_guard lock(state->mutex);
796+
--state->unread_entries;
797+
++state->processed_entries;
798+
return file.file_path.substr(0, file.file_path.find_last_of('-'));
799+
});
800+
if (use_executor) {
801+
rewrite->ScanManifestsWith(executor);
802+
}
803+
ASSERT_THAT(rewrite->Commit(), IsOk());
804+
if (use_executor) {
805+
EXPECT_GT(executor.submit_count(), 0);
806+
}
807+
}
808+
EXPECT_EQ(state->processed_entries, expected_paths.size());
809+
EXPECT_EQ(state->unread_entries, 0);
810+
EXPECT_LE(state->max_unread_entries, 2);
811+
EXPECT_EQ(state->opened_readers, before.size());
812+
EXPECT_EQ(state->open_readers, 0);
813+
EXPECT_LT(state->max_open_readers, before.size());
814+
if (!use_executor) {
815+
EXPECT_EQ(state->max_open_readers, 1);
816+
}
817+
818+
ICEBERG_UNWRAP_OR_FAIL(auto manifests, CurrentManifests());
819+
std::vector<std::string> actual_paths;
820+
for (const auto& manifest : manifests) {
821+
ICEBERG_UNWRAP_OR_FAIL(auto entries, ReadManifestEntries(manifest));
822+
for (const auto& entry : entries) {
823+
EXPECT_EQ(entry.status, ManifestStatus::kExisting);
824+
actual_paths.push_back(entry.data_file->file_path);
825+
}
826+
}
827+
EXPECT_THAT(actual_paths, ::testing::UnorderedElementsAreArray(expected_paths));
828+
}
829+
}
830+
831+
TEST_P(RewriteManifestsTest, CleansWritersAfterStreamingReadFailure) {
832+
ASSERT_THAT(AppendFiles({file_a_, file_b_, file_c_}), IsOk());
833+
ICEBERG_UNWRAP_OR_FAIL(auto before, CurrentManifests());
834+
for (bool use_executor : {false, true}) {
835+
SCOPED_TRACE(use_executor);
836+
auto reads = std::make_shared<ManifestReadState>();
837+
reads->source_paths.insert(before[0].manifest_path);
838+
reads->fail_after_first_batch = true;
839+
auto writes = std::make_shared<ManifestWriteFailureState>();
840+
BindTableWithFileIO(std::make_shared<ManifestWriteFailureFileIO>(
841+
file_io_, table_location_ + "/metadata/", writes));
842+
{
843+
test::ThreadExecutor executor;
844+
ScopedManifestReadTracking tracking(reads);
845+
ICEBERG_UNWRAP_OR_FAIL(auto rewrite, table_->NewRewriteManifests());
846+
rewrite->ClusterBy([](const DataFile&) { return ""; });
847+
if (use_executor) {
848+
rewrite->ScanManifestsWith(executor);
849+
}
850+
auto result = rewrite->Commit();
851+
EXPECT_THAT(result, IsError(ErrorKind::kIOError));
852+
EXPECT_THAT(result, HasErrorMessage("Injected streaming manifest read failure"));
853+
}
854+
EXPECT_EQ(reads->open_readers, 0);
855+
ASSERT_EQ(writes->manifest_outputs.size(), 1U);
856+
EXPECT_TRUE(writes->open_manifest_paths.empty());
857+
EXPECT_TRUE(writes->deleted_open_files.empty());
858+
EXPECT_THAT(writes->deleted_files,
859+
::testing::Contains(writes->manifest_outputs.front()));
860+
EXPECT_FALSE(FileExists(writes->manifest_outputs.front()));
861+
ICEBERG_UNWRAP_OR_FAIL(auto after, CurrentManifests());
862+
EXPECT_EQ(after, before);
863+
}
864+
}
865+
649866
TEST_P(RewriteManifestsTest, RewriteManifestsWithScanExecutor) {
650867
SetSnapshotIdInheritanceEnabled();
651868
ICEBERG_UNWRAP_OR_FAIL(auto new_manifest,

0 commit comments

Comments
 (0)