fix(collection): preserve writing segment on DDL failure (#574)

This commit is contained in:
egolearner 2026-07-16 10:54:56 +08:00 committed by GitHub
parent 8693cf9de7
commit 23538ab876
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 323 additions and 130 deletions

View File

@ -160,6 +160,11 @@ class CollectionImpl : public Collection {
Status switch_to_new_segment_for_writing(
const CollectionSchema::Ptr &schema = nullptr);
Status commit_schema_change_with_new_writing_segment(
const CollectionSchema::Ptr &new_schema,
const Segment::Ptr &old_writing_segment, const Version &old_version,
Version *new_version, uint64_t writing_min_doc_id);
Result<WriteResults> write_impl(std::vector<Doc> &docs, WriteMode mode);
std::vector<Segment::Ptr> get_all_segments() const;
@ -493,54 +498,18 @@ Status CollectionImpl::CreateIndex(const std::string &column_name,
// forbidden writing until index is ready
std::lock_guard write_lock(write_mtx_);
Version new_version = version_manager_->get_current_version();
if (writing_segment_->doc_count() > 0) {
s = writing_segment_->dump();
s = switch_to_new_segment_for_writing();
CHECK_RETURN_STATUS(s);
s = segment_manager_->add_segment(writing_segment_);
CHECK_RETURN_STATUS(s);
auto seg_options =
SegmentOptions{false, options_.enable_mmap_, options_.max_buffer_size_};
auto new_segment = Segment::CreateAndOpen(
path_, *new_schema, allocate_segment_id(),
writing_segment_->meta()->max_doc_id() + 1, id_map_, delete_store_,
version_manager_, seg_options);
if (!new_segment) {
return new_segment.error();
}
s = new_version.add_persisted_segment_meta(writing_segment_->meta());
CHECK_RETURN_STATUS(s);
writing_segment_ = new_segment.value();
new_version.set_next_segment_id(segment_id_allocator_.load());
} else {
// TODO: allocate new segment id and clear current writing segment at last
// recreate writing segment
s = writing_segment_->destroy();
CHECK_RETURN_STATUS(s);
auto id = writing_segment_->id();
auto min_doc_id = writing_segment_->meta()->min_doc_id();
writing_segment_.reset();
SegmentOptions seg_options;
seg_options.enable_mmap_ = options_.enable_mmap_;
seg_options.max_buffer_size_ = options_.max_buffer_size_;
seg_options.read_only_ = options_.read_only_;
auto writing_segment =
Segment::CreateAndOpen(path_, *new_schema, id, min_doc_id, id_map_,
delete_store_, version_manager_, seg_options);
if (!writing_segment) {
return writing_segment.error();
}
writing_segment_ = writing_segment.value();
}
new_version.reset_writing_segment_meta(writing_segment_->meta());
// get_all_segment will return writing segment if it has docs
auto old_writing_segment = writing_segment_;
Version old_version = version_manager_->get_current_version();
Version new_version = old_version;
auto writing_min_doc_id = old_writing_segment->meta()->min_doc_id();
// DDL tasks only run on persisted segments. Non-empty writing segment has
// already been switched to a persisted segment above.
auto persist_segments = get_all_persist_segments();
bool is_vector_field = field->is_vector_field();
@ -562,23 +531,11 @@ Status CollectionImpl::CreateIndex(const std::string &column_name,
"] is not supported");
}
if (tasks.empty()) {
new_version.set_schema(*new_schema);
s = version_manager_->apply(new_version);
if (!tasks.empty()) {
s = execute_tasks(tasks);
CHECK_RETURN_STATUS(s);
// persist manifest
s = version_manager_->flush();
CHECK_RETURN_STATUS(s);
schema_ = new_schema;
return Status::OK();
}
s = execute_tasks(tasks);
CHECK_RETURN_STATUS(s);
new_version.set_schema(*new_schema);
for (auto &task : tasks) {
@ -600,12 +557,9 @@ Status CollectionImpl::CreateIndex(const std::string &column_name,
CHECK_RETURN_STATUS(s);
}
// 2. update version
s = version_manager_->apply(new_version);
CHECK_RETURN_STATUS(s);
// 3. persist version
s = version_manager_->flush();
s = commit_schema_change_with_new_writing_segment(
new_schema, old_writing_segment, old_version, &new_version,
writing_min_doc_id);
CHECK_RETURN_STATUS(s);
// 4. remove old segments or block
@ -632,8 +586,6 @@ Status CollectionImpl::CreateIndex(const std::string &column_name,
CHECK_RETURN_STATUS(s);
}
schema_ = new_schema;
return Status::OK();
}
@ -719,52 +671,15 @@ Status CollectionImpl::DropIndex(const std::string &column_name) {
// forbidden writing until index is ready
std::lock_guard write_lock(write_mtx_);
Version new_version = version_manager_->get_current_version();
if (writing_segment_->doc_count() > 0) {
s = writing_segment_->dump();
s = switch_to_new_segment_for_writing();
CHECK_RETURN_STATUS(s);
s = segment_manager_->add_segment(writing_segment_);
CHECK_RETURN_STATUS(s);
auto new_segment =
Segment::CreateAndOpen(path_, *new_schema, allocate_segment_id(),
writing_segment_->meta()->max_doc_id() + 1,
id_map_, delete_store_, version_manager_,
SegmentOptions{false, options_.enable_mmap_,
options_.max_buffer_size_});
if (!new_segment) {
return new_segment.error();
}
s = new_version.add_persisted_segment_meta(writing_segment_->meta());
CHECK_RETURN_STATUS(s);
writing_segment_ = new_segment.value();
new_version.set_next_segment_id(segment_id_allocator_.load());
} else {
// recreate writing segment
s = writing_segment_->destroy();
CHECK_RETURN_STATUS(s);
auto id = writing_segment_->id();
auto min_doc_id = writing_segment_->meta()->min_doc_id();
writing_segment_.reset();
SegmentOptions seg_options;
seg_options.enable_mmap_ = options_.enable_mmap_;
seg_options.max_buffer_size_ = options_.max_buffer_size_;
seg_options.read_only_ = options_.read_only_;
auto writing_segment =
Segment::CreateAndOpen(path_, *new_schema, id, min_doc_id, id_map_,
delete_store_, version_manager_, seg_options);
if (!writing_segment) {
return writing_segment.error();
}
writing_segment_ = writing_segment.value();
}
new_version.reset_writing_segment_meta(writing_segment_->meta());
auto old_writing_segment = writing_segment_;
Version old_version = version_manager_->get_current_version();
Version new_version = old_version;
auto writing_min_doc_id = old_writing_segment->meta()->min_doc_id();
auto persist_segments = get_all_persist_segments();
@ -784,23 +699,11 @@ Status CollectionImpl::DropIndex(const std::string &column_name) {
"] on column[", column_name, "] is not supported");
}
if (tasks.empty()) {
new_version.set_schema(*new_schema);
s = version_manager_->apply(new_version);
if (!tasks.empty()) {
s = execute_tasks(tasks);
CHECK_RETURN_STATUS(s);
// persist manifest
s = version_manager_->flush();
CHECK_RETURN_STATUS(s);
schema_ = new_schema;
return Status::OK();
}
s = execute_tasks(tasks);
CHECK_RETURN_STATUS(s);
new_version.set_schema(*new_schema);
for (auto &task : tasks) {
@ -822,11 +725,9 @@ Status CollectionImpl::DropIndex(const std::string &column_name) {
CHECK_RETURN_STATUS(s);
}
s = version_manager_->apply(new_version);
CHECK_RETURN_STATUS(s);
// persist manifest
s = version_manager_->flush();
s = commit_schema_change_with_new_writing_segment(
new_schema, old_writing_segment, old_version, &new_version,
writing_min_doc_id);
CHECK_RETURN_STATUS(s);
// 4. remove old segments or block
@ -852,8 +753,6 @@ Status CollectionImpl::DropIndex(const std::string &column_name) {
CHECK_RETURN_STATUS(s);
}
schema_ = new_schema;
return Status::OK();
}
@ -1574,6 +1473,47 @@ bool CollectionImpl::need_switch_to_new_segment() const {
return writing_segment_->doc_count() >= schema_->max_doc_count_per_segment();
}
Status CollectionImpl::commit_schema_change_with_new_writing_segment(
const CollectionSchema::Ptr &new_schema,
const Segment::Ptr &old_writing_segment, const Version &old_version,
Version *new_version, uint64_t writing_min_doc_id) {
if (new_version == nullptr) {
return Status::InvalidArgument("new_version is null");
}
auto seg_options =
SegmentOptions{false, options_.enable_mmap_, options_.max_buffer_size_};
auto new_writing_segment = Segment::CreateAndOpen(
path_, *new_schema, allocate_segment_id(), writing_min_doc_id, id_map_,
delete_store_, version_manager_, seg_options);
if (!new_writing_segment) {
return new_writing_segment.error();
}
new_version->reset_writing_segment_meta(new_writing_segment.value()->meta());
new_version->set_next_segment_id(segment_id_allocator_.load());
auto s = version_manager_->apply(*new_version);
if (!s.ok()) {
new_writing_segment.value()->destroy();
return s;
}
s = version_manager_->flush();
if (!s.ok()) {
new_writing_segment.value()->destroy();
auto rollback_status = version_manager_->apply(old_version);
CHECK_RETURN_STATUS(rollback_status);
return s;
}
schema_ = new_schema;
writing_segment_ = new_writing_segment.value();
s = old_writing_segment->destroy();
CHECK_RETURN_STATUS(s);
return Status::OK();
}
Status CollectionImpl::switch_to_new_segment_for_writing(
const CollectionSchema::Ptr &schema) {
auto s = writing_segment_->dump();

View File

@ -17,6 +17,7 @@
#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <memory>
#include <mutex>
@ -60,6 +61,73 @@ class CollectionTest : public ::testing::Test {
}
};
class DirectoryWriteBlockerForTest {
public:
explicit DirectoryWriteBlockerForTest(const std::string &dir_path)
: dir_path_(dir_path) {
namespace fs = std::filesystem;
std::error_code ec;
auto status = fs::status(dir_path_, ec);
if (ec) {
skip_reason_ = "Failed to stat directory: " + ec.message();
return;
}
original_perms_ = status.permissions();
fs::permissions(dir_path_,
fs::perms::owner_read | fs::perms::owner_exec |
fs::perms::group_read | fs::perms::group_exec |
fs::perms::others_read | fs::perms::others_exec,
fs::perm_options::replace, ec);
if (ec) {
skip_reason_ = "Failed to make directory read-only: " + ec.message();
return;
}
enabled_ = true;
auto probe_path = (fs::path(dir_path_) / ".zvec_permission_probe").string();
{
std::ofstream probe_file(probe_path, std::ios::out | std::ios::trunc);
if (probe_file.is_open()) {
probe_file << "probe";
probe_file.close();
fs::remove(probe_path, ec);
Restore();
skip_reason_ = "Directory permissions do not block writes";
return;
}
}
}
~DirectoryWriteBlockerForTest() {
Restore();
}
bool enabled() const {
return enabled_;
}
const std::string &skip_reason() const {
return skip_reason_;
}
void Restore() {
if (enabled_) {
std::error_code ec;
std::filesystem::permissions(dir_path_, original_perms_,
std::filesystem::perm_options::replace, ec);
enabled_ = false;
}
}
private:
std::string dir_path_;
std::filesystem::perms original_perms_{std::filesystem::perms::unknown};
bool enabled_{false};
std::string skip_reason_;
};
TEST_F(CollectionTest, Feature_CreateAndOpen_General) {
auto func = [&](bool enable_mmap) {
CollectionOptions options;
@ -2600,6 +2668,123 @@ TEST_F(CollectionTest, Feature_DropIndex_Scalar) {
}
}
TEST_F(CollectionTest,
Feature_DropIndex_Scalar_FailureKeepsPersistedOldSchema) {
#ifdef __ANDROID__
GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink "
"support (needed by RocksDB checkpoint)";
#endif
FileHelper::RemoveDirectory(col_path);
int doc_count = 1;
auto schema = TestHelper::CreateSchemaWithScalarIndex(false, true);
auto options = CollectionOptions{false, true, 64 * 1024 * 1024};
auto collection = TestHelper::CreateCollectionWithDoc(
col_path, *schema, options, 0, doc_count, false);
ASSERT_TRUE(collection->Optimize().ok());
collection.reset();
auto reopen_result = Collection::Open(col_path, options);
ASSERT_TRUE(reopen_result.has_value()) << reopen_result.error().message();
collection = std::move(reopen_result.value());
auto segment_path = FileHelper::MakeSegmentPath(col_path, 0);
DirectoryWriteBlockerForTest write_blocker(segment_path);
if (!write_blocker.enabled()) {
GTEST_SKIP() << write_blocker.skip_reason();
}
auto s = collection->DropIndex("int32");
write_blocker.Restore();
ASSERT_FALSE(s.ok());
collection.reset();
reopen_result = Collection::Open(col_path, options);
ASSERT_TRUE(reopen_result.has_value()) << reopen_result.error().message();
collection = std::move(reopen_result.value());
auto schema_after_drop = collection->Schema();
ASSERT_TRUE(schema_after_drop.has_value())
<< schema_after_drop.error().message();
ASSERT_NE(schema_after_drop.value().get_field("string")->index_params(),
nullptr);
ASSERT_NE(schema_after_drop.value().get_field("int32")->index_params(),
nullptr);
ASSERT_EQ(collection->Stats().value().doc_count, doc_count);
auto expect_doc = TestHelper::CreateDoc(0, *schema);
auto result = collection->Fetch({expect_doc.pk()});
ASSERT_TRUE(result.has_value());
ASSERT_EQ(result.value().size(), 1);
ASSERT_EQ(result.value().count(expect_doc.pk()), 1);
auto doc = result.value()[expect_doc.pk()];
ASSERT_NE(doc, nullptr);
ASSERT_EQ(*doc, expect_doc);
collection.reset();
FileHelper::RemoveDirectory(col_path);
}
TEST_F(CollectionTest, Feature_IndexDDL_WritingSegmentReopen) {
auto run = [&](bool create_index, bool insert_before_ddl) {
FileHelper::RemoveDirectory(col_path);
auto index_params = std::make_shared<InvertIndexParams>();
auto schema = std::make_shared<CollectionSchema>(
create_index ? "create_index_writing" : "drop_index_writing");
schema->add_field(
std::make_shared<FieldSchema>("name", DataType::STRING, false,
create_index ? nullptr : index_params));
auto options = CollectionOptions{false, true, 64 * 1024 * 1024};
auto collection_res = Collection::CreateAndOpen(col_path, *schema, options);
ASSERT_TRUE(collection_res.has_value()) << collection_res.error().message();
auto collection = std::move(collection_res.value());
if (insert_before_ddl) {
Doc doc;
doc.set_pk("pk0");
doc.set<std::string>("name", "hello world");
std::vector<Doc> docs{doc};
ASSERT_TRUE(collection->Insert(docs).has_value());
}
auto s = create_index ? collection->CreateIndex("name", index_params)
: collection->DropIndex("name");
ASSERT_TRUE(s.ok()) << s.message();
collection.reset();
auto reopen_result = Collection::Open(col_path, options);
ASSERT_TRUE(reopen_result.has_value()) << reopen_result.error().message();
collection = std::move(reopen_result.value());
auto schema_after_ddl = collection->Schema();
ASSERT_TRUE(schema_after_ddl.has_value())
<< schema_after_ddl.error().message();
bool has_index =
schema_after_ddl.value().get_field("name")->index_params() != nullptr;
ASSERT_EQ(has_index, create_index);
ASSERT_EQ(collection->Stats().value().doc_count,
insert_before_ddl ? 1u : 0u);
if (insert_before_ddl) {
auto result = collection->Fetch({"pk0"});
ASSERT_TRUE(result.has_value()) << result.error().message();
ASSERT_EQ(result.value().size(), 1u);
ASSERT_EQ(result.value()["pk0"]->get<std::string>("name").value(),
"hello world");
}
collection.reset();
FileHelper::RemoveDirectory(col_path);
};
run(true, false);
run(false, false);
run(true, true);
run(false, true);
}
TEST_F(CollectionTest, Feature_DropIndex_AfterCreate) {
auto func = [&](std::string field_name, bool enable_optimize) {
FileHelper::RemoveDirectory(col_path);
@ -6101,6 +6286,74 @@ TEST_F(CollectionTest, Feature_NoVectorCollection_FtsReopenThenInsert) {
FileHelper::RemoveDirectory(col_path);
}
// Force the persisted segment FTS snapshot task to fail before schema commit,
// then verify reopen still sees the old schema and data.
TEST_F(CollectionTest, Feature_DropFtsIndex_FailureKeepsPersistedOldSchema) {
#ifdef __ANDROID__
GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink "
"support (needed by RocksDB checkpoint)";
#endif
FileHelper::RemoveDirectory(col_path);
auto schema = std::make_shared<CollectionSchema>("fts_drop_reopen");
schema->add_field(std::make_shared<FieldSchema>("title", DataType::STRING));
schema->add_field(std::make_shared<FieldSchema>(
"content", DataType::STRING, false, std::make_shared<FtsIndexParams>()));
schema->add_field(
std::make_shared<FieldSchema>("other_content", DataType::STRING, false,
std::make_shared<FtsIndexParams>()));
CollectionOptions options{false, true};
auto col_res = Collection::CreateAndOpen(col_path, *schema, options);
ASSERT_TRUE(col_res.has_value()) << col_res.error().message();
auto col = std::move(col_res.value());
Doc doc;
doc.set_pk("pk0");
doc.set<std::string>("title", "title");
doc.set<std::string>("content", "hello world");
doc.set<std::string>("other_content", "hello other");
std::vector<Doc> docs{doc};
ASSERT_TRUE(col->Insert(docs).has_value());
ASSERT_TRUE(col->Optimize().ok());
col.reset();
auto reopen_res = Collection::Open(col_path, options);
ASSERT_TRUE(reopen_res.has_value()) << reopen_res.error().message();
col = std::move(reopen_res.value());
auto segment_path = FileHelper::MakeSegmentPath(col_path, 0);
DirectoryWriteBlockerForTest write_blocker(segment_path);
if (!write_blocker.enabled()) {
GTEST_SKIP() << write_blocker.skip_reason();
}
auto s = col->DropIndex("content");
write_blocker.Restore();
ASSERT_FALSE(s.ok());
col.reset();
reopen_res = Collection::Open(col_path, options);
ASSERT_TRUE(reopen_res.has_value()) << reopen_res.error().message();
col = std::move(reopen_res.value());
auto schema_after_drop = col->Schema();
ASSERT_TRUE(schema_after_drop.has_value())
<< schema_after_drop.error().message();
ASSERT_NE(schema_after_drop.value().get_field("content")->index_params(),
nullptr);
ASSERT_NE(
schema_after_drop.value().get_field("other_content")->index_params(),
nullptr);
ASSERT_EQ(col->Stats().value().doc_count, 1u);
auto fetched = col->Fetch({"pk0"});
ASSERT_TRUE(fetched.has_value()) << fetched.error().message();
ASSERT_EQ(fetched.value().size(), 1u);
col.reset();
FileHelper::RemoveDirectory(col_path);
}
// Dynamic CreateIndex/DropIndex for FTS: create an FTS index on a STRING column
// that already has data, verify queries hit, then drop the index and verify FTS
// is no longer available. Also covers reopen persistence.