From 331933c75a2240eb72b070a0ae87b2e330cb3f42 Mon Sep 17 00:00:00 2001 From: Qinren Zhou Date: Wed, 11 Mar 2026 15:55:13 +0800 Subject: [PATCH] fix: clean up crash residue (#208) --- src/db/index/common/doc.cc | 63 +++- src/db/index/segment/segment.cc | 28 +- src/include/zvec/db/doc.h | 2 + tests/db/CMakeLists.txt | 1 + tests/db/crash_recovery/CMakeLists.txt | 67 ++++ tests/db/crash_recovery/data_generator.cc | 217 +++++++++++ tests/db/crash_recovery/utility.h | 152 ++++++++ .../db/crash_recovery/write_recovery_test.cc | 357 ++++++++++++++++++ 8 files changed, 865 insertions(+), 22 deletions(-) create mode 100644 tests/db/crash_recovery/CMakeLists.txt create mode 100644 tests/db/crash_recovery/data_generator.cc create mode 100644 tests/db/crash_recovery/utility.h create mode 100644 tests/db/crash_recovery/write_recovery_test.cc diff --git a/src/db/index/common/doc.cc b/src/db/index/common/doc.cc index dad9bbd..6d411bf 100644 --- a/src/db/index/common/doc.cc +++ b/src/db/index/common/doc.cc @@ -11,6 +11,7 @@ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include #include @@ -1108,6 +1109,52 @@ std::string Doc::to_detail_string() const { return oss.str(); } +struct Doc::ValueEqual { + template + bool operator()(const T &, const U &) const { + return false; + } + + template + bool operator()(const T &a, const T &b) const { + return a == b; + } + + bool operator()(float a, float b) const { + return std::fabs(a - b) < 1e-6f; + } + + bool operator()(double a, double b) const { + return std::fabs(a - b) < 1e-9; + } + + bool operator()(const std::vector &a, + const std::vector &b) const { + if (a.size() != b.size()) return false; + for (size_t i = 0; i < a.size(); ++i) + if (std::fabs(static_cast(a[i]) - static_cast(b[i])) >= + 1e-3f) + return false; + return true; + } + + bool operator()(const std::vector &a, + const std::vector &b) const { + if (a.size() != b.size()) return false; + for (size_t i = 0; i < a.size(); ++i) + if (std::fabs(a[i] - b[i]) >= 1e-6f) return false; + return true; + } + + bool operator()(const std::vector &a, + const std::vector &b) const { + if (a.size() != b.size()) return false; + for (size_t i = 0; i < a.size(); ++i) + if (std::fabs(a[i] - b[i]) >= 1e-9) return false; + return true; + } +}; + bool Doc::operator==(const Doc &other) const { // Compare basic fields if (pk_ != other.pk_) { @@ -1135,21 +1182,7 @@ bool Doc::operator==(const Doc &other) const { } // Use visitor to compare the actual values - bool values_equal = std::visit( - [](const auto &lhs, const auto &rhs) -> bool { - if constexpr (std::is_same_v, - std::decay_t>) { - return lhs == rhs; - } else { - // This should not happen due to the index check above - return false; - } - }, - field_value, it->second); - - if (!values_equal) { - return false; - } + if (!std::visit(ValueEqual{}, field_value, it->second)) return false; } return true; diff --git a/src/db/index/segment/segment.cc b/src/db/index/segment/segment.cc index 517215a..2d03cd7 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -3939,6 +3939,14 @@ VectorColumnIndexer::Ptr SegmentImpl::create_vector_indexer( memory_vector_block_ids_[field_name] = block_id; } + if (FileHelper::FileExists(index_file_path)) { + LOG_WARN( + "Index file[%s] already exists (possible crash residue); cleaning and " + "overwriting.", + index_file_path.c_str()); + FileHelper::RemoveFile(index_file_path); + } + auto vector_indexer = std::make_shared(index_file_path, field); vector_column_params::ReadOptions options{true, true}; @@ -3958,6 +3966,13 @@ Status SegmentImpl::init_memory_components() { // create and open memory forward block auto mem_path = FileHelper::MakeForwardBlockPath(seg_path_, mem_block.id_, !options_.enable_mmap_); + if (FileHelper::FileExists(mem_path)) { + LOG_WARN( + "ForwardBlock file[%s] already exists (possible crash residue); " + "cleaning and overwriting.", + mem_path.c_str()); + FileHelper::RemoveFile(mem_path); + } memory_store_ = std::make_shared( collection_schema_, mem_path, options_.enable_mmap_ ? FileFormat::IPC : FileFormat::PARQUET, @@ -4104,18 +4119,17 @@ Status SegmentImpl::recover() { } const auto added_docs = recovered_doc_count[0] + // INSERT - recovered_doc_count[1] + // UPDATE - recovered_doc_count[2]; // UPSERT + recovered_doc_count[1] + // UPSERT + recovered_doc_count[2]; // UPDATE mem_block.max_doc_id_ += added_docs; LOG_INFO( - "Recover from wal finished. total_recovered_doc_count[%zu] " - "insert[%zu] update[%zu] upsert[%zu] " - "delete[%zu] path[%s]", + "Recover from wal finished. total_recovered_doc_count[%zu] insert[%zu] " + "upsert[%zu] update[%zu] delete[%zu] path[%s]", (size_t)total_recovered_doc_count, (size_t)recovered_doc_count[0], // INSERT - (size_t)recovered_doc_count[1], // UPDATE - (size_t)recovered_doc_count[2], // UPSERT + (size_t)recovered_doc_count[1], // UPSERT + (size_t)recovered_doc_count[2], // UPDATE (size_t)recovered_doc_count[3], // DELETE wal_file_path.c_str()); diff --git a/src/include/zvec/db/doc.h b/src/include/zvec/db/doc.h index 5f927fa..fa05605 100644 --- a/src/include/zvec/db/doc.h +++ b/src/include/zvec/db/doc.h @@ -294,6 +294,8 @@ class Doc { static void read_from_buffer(const uint8_t *&data, void *dest, size_t size); + struct ValueEqual; + private: std::string pk_; float score_{0.0f}; diff --git a/tests/db/CMakeLists.txt b/tests/db/CMakeLists.txt index 8de3089..612ee15 100644 --- a/tests/db/CMakeLists.txt +++ b/tests/db/CMakeLists.txt @@ -2,6 +2,7 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) include(${PROJECT_ROOT_DIR}/cmake/option.cmake) cc_directory(common) +cc_directories(crash_recovery) cc_directory(sqlengine) cc_directories(index) diff --git a/tests/db/crash_recovery/CMakeLists.txt b/tests/db/crash_recovery/CMakeLists.txt new file mode 100644 index 0000000..296b8c5 --- /dev/null +++ b/tests/db/crash_recovery/CMakeLists.txt @@ -0,0 +1,67 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +if(APPLE) + set(APPLE_FRAMEWORK_LIBS + -framework CoreFoundation + -framework CoreGraphics + -framework CoreData + -framework CoreText + -framework Security + -framework Foundation + -Wl,-U,_MallocExtension_ReleaseFreeMemory + -Wl,-U,_ProfilerStart + -Wl,-U,_ProfilerStop + -Wl,-U,_RegisterThriftProtocol + ) +endif() + + +# Build data_generator executable +cc_binary( + NAME data_generator + LIBS zvec_db + zvec_proto + core_knn_flat + core_knn_flat_sparse + core_knn_hnsw + core_knn_hnsw_sparse + core_knn_ivf + core_mix_reducer + core_metric + core_utility + core_quantizer + ${CMAKE_THREAD_LIBS_INIT} + ${CMAKE_DL_LIBS} + SRCS data_generator.cc + INCS .. ../../src + LDFLAGS ${APPLE_FRAMEWORK_LIBS} +) + + +# Build test executables +file(GLOB ALL_TEST_SRCS *_test.cc) +foreach(CC_SRCS ${ALL_TEST_SRCS}) + get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE) + cc_gmock( + NAME ${CC_TARGET} STRICT + LIBS zvec_db + zvec_proto + core_knn_flat + core_knn_flat_sparse + core_knn_hnsw + core_knn_hnsw_sparse + core_knn_ivf + core_mix_reducer + core_metric + core_utility + core_quantizer + ${CMAKE_THREAD_LIBS_INIT} + ${CMAKE_DL_LIBS} + SRCS ${CC_SRCS} + INCS .. ../../src + LDFLAGS ${APPLE_FRAMEWORK_LIBS} + ) + add_dependencies(${CC_TARGET} data_generator) + cc_test_suite(zvec_crash_recovery ${CC_TARGET}) +endforeach() diff --git a/tests/db/crash_recovery/data_generator.cc b/tests/db/crash_recovery/data_generator.cc new file mode 100644 index 0000000..5754247 --- /dev/null +++ b/tests/db/crash_recovery/data_generator.cc @@ -0,0 +1,217 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + + +#include +#include +#include +#include +#include "zvec/ailego/logger/logger.h" +#include "utility.h" + + +constexpr int kBatchSize = 20; +constexpr int kBatchDelayMs = 10; + + +struct Config { + std::string path; + int start_id = 0; + int end_id = 0; + std::string operation; // "insert", "upsert", "update", "delete" + int version = 999999; +}; + + +bool ParseArgs(int argc, char **argv, Config &config) { + for (int i = 1; i < argc; i++) { + std::string arg = argv[i]; + + if (arg == "--path" && i + 1 < argc) { + config.path = argv[++i]; + } else if (arg == "--start" && i + 1 < argc) { + config.start_id = std::stoi(argv[++i]); + } else if (arg == "--end" && i + 1 < argc) { + config.end_id = std::stoi(argv[++i]); + } else if (arg == "--op" && i + 1 < argc) { + config.operation = argv[++i]; + } else if (arg == "--version" && i + 1 < argc) { + config.version = std::stoi(argv[++i]); + } else if (arg == "--help" || arg == "-h") { + return false; + } + } + + // Validate required arguments + if (config.path.empty() || config.operation.empty() || + config.start_id >= config.end_id || config.version == 999999) { + return false; + } + + // Validate operation + if (config.operation != "insert" && config.operation != "upsert" && + config.operation != "update" && config.operation != "delete") { + std::cerr << "Error: Invalid operation '" << config.operation + << "'. Must be 'insert', 'upsert', 'update', or 'delete'." + << std::endl; + return false; + } + + return true; +} + + +void PrintUsage(const char *program) { + std::cout << "Usage: " << program + << " --path --start --end " + "--op " + << std::endl; + std::cout << std::endl; + std::cout << "Arguments:" << std::endl; + std::cout << " --path Path to the collection (required)" << std::endl; + std::cout << " --start Starting document ID (inclusive, required)" + << std::endl; + std::cout << " --end Ending document ID (exclusive, required)" + << std::endl; + std::cout + << " --op Operation: insert, upsert, update, or delete (required)" + << std::endl; + std::cout << " --version Operation: version (required)" << std::endl; + std::cout << std::endl; + std::cout << "Examples:" << std::endl; + std::cout << " # Insert 1000 documents (pk_0 to pk_999)" << std::endl; + std::cout << " " << program + << " --path ./test_db --start 0 --end 1000 --op insert --version 0" + << std::endl; + std::cout << std::endl; + std::cout << " # Update documents 1000-1999" << std::endl; + std::cout + << " " << program + << " --path ./test_db --start 1000 --end 2000 --op update --version 1" + << std::endl; + std::cout << std::endl; + std::cout << " # Upsert documents 0-499" << std::endl; + std::cout << " " << program + << " --path ./test_db --start 0 --end 500 --op upsert --version 2" + << std::endl; +} + + +int main(int argc, char **argv) { + Config config; + + // Parse arguments + if (!ParseArgs(argc, argv, config)) { + PrintUsage(argv[0]); + return 1; + } + + try { + std::filesystem::path cwd = std::filesystem::current_path(); + std::cout << "[data_generator] Current Working Directory: " << cwd.string() + << std::endl; + } catch (const std::filesystem::filesystem_error &e) { + std::cout + << "[data_generator] Failed to get the current working directory: " + << e.what() << std::endl; + } + + std::cout << "Configuration:" << std::endl; + std::cout << " Path: " << config.path << std::endl; + std::cout << " Range: [" << config.start_id << ", " << config.end_id + << ")" << std::endl; + std::cout << " Operation: " << config.operation << std::endl; + std::cout << " BatchSize: " << kBatchSize << std::endl; + std::cout << " BatchDelay: " << kBatchDelayMs << "ms" << std::endl; + std::cout << std::endl; + + auto result = + zvec::Collection::Open(config.path, zvec::CollectionOptions{false, true}); + if (!result) { + LOG_ERROR("Failed to open collection[%s]: %s", config.path.c_str(), + result.error().c_str()); + return -1; + } + + auto collection = result.value(); + LOG_INFO("Collection[%s] opened successfully", config.path.c_str()); + + // Process documents in batches + int total_docs = config.end_id - config.start_id; + int processed = 0; + int batch_num = 0; + int next_progress_threshold = total_docs / 10; // 10% increments + int progress_percent = 0; + + while (config.start_id < config.end_id) { + int batch_end = std::min(config.start_id + kBatchSize, config.end_id); + int batch_count = batch_end - config.start_id; + + std::vector docs; + docs.reserve(batch_count); + for (uint64_t i = config.start_id; i < batch_end; i++) { + docs.push_back(zvec::CreateTestDoc(i, config.version)); + } + + zvec::Result results; + if (config.operation == "insert") { + results = collection->Insert(docs); + } else if (config.operation == "upsert") { + results = collection->Upsert(docs); + } else if (config.operation == "update") { + results = collection->Update(docs); + } else if (config.operation == "delete") { + std::vector pks{}; + for (const auto &doc : docs) { + pks.emplace_back(doc.pk()); + } + results = collection->Delete(pks); + } + if (!results) { + LOG_ERROR("Failed to perform operation[%s], reason: %s", + config.operation.c_str(), results.error().message().c_str()); + return 1; + } + for (auto &s : results.value()) { + if (!s.ok()) { + LOG_ERROR("Failed to perform operation[%s], reason: %s", + config.operation.c_str(), s.message().c_str()); + return 1; + } + } + + processed += batch_count; + config.start_id = batch_end; + batch_num++; + + // Print progress every 10% + if (processed >= next_progress_threshold) { + progress_percent++; + LOG_INFO("Progress: %d (%d/%d documents)", progress_percent * 10, + processed, total_docs); + next_progress_threshold = (progress_percent + 1) * total_docs / 10; + } + + // Sleep between batches + if (config.start_id < config.end_id) { + std::this_thread::sleep_for(std::chrono::milliseconds(kBatchDelayMs)); + } + } + + std::cout << std::endl; + std::cout << "Success! Processed " << processed << " documents in " + << batch_num << " batches." << std::endl; + + return 0; +} diff --git a/tests/db/crash_recovery/utility.h b/tests/db/crash_recovery/utility.h new file mode 100644 index 0000000..36768b2 --- /dev/null +++ b/tests/db/crash_recovery/utility.h @@ -0,0 +1,152 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + + +#pragma once + + +#include +#include + + +namespace zvec { + +/** + * @brief Create a test schema with deterministic field definitions. + * + * @param name The collection name (default: "crash_recovery_test") + * @return CollectionSchema::Ptr The test schema + */ +inline CollectionSchema::Ptr CreateTestSchema( + const std::string &name = "crash_recovery_test") { + auto schema = std::make_shared(name); + schema->set_max_doc_count_per_segment(2000); + + schema->add_field( + std::make_shared("int32_field", DataType::INT32, false)); + schema->add_field( + std::make_shared("int64_field", DataType::INT64, true)); + schema->add_field( + std::make_shared("float_field", DataType::FLOAT, true)); + schema->add_field( + std::make_shared("string_field", DataType::STRING, false)); + schema->add_field( + std::make_shared("bool_field", DataType::BOOL, false)); + schema->add_field(std::make_shared("array_int32_field", + DataType::ARRAY_INT32, true)); + schema->add_field(std::make_shared( + "array_string_field", DataType::ARRAY_STRING, false)); + schema->add_field(std::make_shared( + "dense_fp32_field", DataType::VECTOR_FP32, 128, false, + std::make_shared(MetricType::COSINE))); + schema->add_field(std::make_shared( + "sparse_fp32_field", DataType::SPARSE_VECTOR_FP32, 0, false, + std::make_shared(MetricType::IP))); + + return schema; +} + + +/** + * @brief Create a test document with deterministic values based on doc_id. + * + * Document pattern: + * - pk: "pk_{doc_id}" + * - int32_field: doc_id (cast to int32) + * - int64_field: doc_id, null if doc_id % 60 == 0 + * - float_field: doc_id / 1000.0, null if doc_id % 70 == 0 + * - string_field: "{version}_{doc_id}" + * - bool_field: doc_id % 2 == 0 or flipped if version % 2 !=0 + * - array_int32_field: [doc_id, doc_id+1, doc_id+2], null if doc_id % 100 == 0 + * - array_string_field: ["str_{version}_0", ...] + * - dense_fp32_field: vector where dense[i] = (doc_id + i) / 1000.0f + * - sparse_fp32_field: sparse vector with indices [0, 10, ...] + * + * @param doc_id The document ID (determines all field values) + * @param version The version of the document + * @return Doc The created document + */ +inline Doc CreateTestDoc(uint64_t doc_id, int version) { + Doc doc; + + // Set primary key + std::string pk = "pk_" + std::to_string(doc_id); + doc.set_pk(pk); + + // Set scalar fields + doc.set("int32_field", static_cast(doc_id)); + + // int64_field: nullable, null if doc_id % 60 == 0 + if (doc_id % 60 != 0) { + doc.set("int64_field", static_cast(doc_id)); + } + + // float_field: nullable, null if doc_id % 70 == 0 + if (doc_id % 70 != 0) { + doc.set("float_field", static_cast(doc_id) / 1000.0f); + } + + // string_field: "value_{id}" or "updated_value_{id}" + std::string string_value = + std::to_string(version) + "_" + std::to_string(doc_id); + doc.set("string_field", string_value); + + // bool_field: alternating based on doc_id, flipped if updated + bool bool_value = (doc_id % 2 == 0); + if (version % 2 != 0) { + bool_value = !bool_value; + } + doc.set("bool_field", bool_value); + + // array_int32_field: nullable, null if doc_id % 100 == 0 + if (doc_id % 100 != 0) { + std::vector array_int32; + for (int i = 0; i < 3; i++) { + array_int32.push_back(static_cast(doc_id + i)); + } + doc.set>("array_int32_field", array_int32); + } + + // array_string_field: ["str_0", "str_1", ...] or ["updated_str_0", ...] + std::vector array_string; + size_t array_size = doc_id % 5 + 1; // 1 to 5 elements + for (size_t i = 0; i < array_size; i++) { + array_string.push_back("str_" + std::to_string(version) + "_" + + std::to_string(i)); + } + doc.set>("array_string_field", array_string); + + // dense_fp32_field: deterministic pattern + std::vector dense(128); + for (int i = 0; i < 128; i++) { + dense[i] = static_cast(doc_id + i) / 1000.0f; + } + doc.set>("dense_fp32_field", dense); + + // sparse_fp32_field: sparse vector with indices [0, 10, 20, ..., 100] + // Values based on doc_id: value = (doc_id + index) / 1000.0 + std::vector sparse_indices; + std::vector sparse_values; + for (uint32_t idx = 0; idx <= 100; idx += 10) { + sparse_indices.push_back(idx); + sparse_values.push_back(static_cast(doc_id + idx) / 1000.0f); + } + doc.set, std::vector>>( + "sparse_fp32_field", std::make_pair(sparse_indices, sparse_values)); + + return doc; +} + + +} // namespace zvec diff --git a/tests/db/crash_recovery/write_recovery_test.cc b/tests/db/crash_recovery/write_recovery_test.cc new file mode 100644 index 0000000..1f53a5f --- /dev/null +++ b/tests/db/crash_recovery/write_recovery_test.cc @@ -0,0 +1,357 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + + +#include +#include +#include +#include +#include +#include +#include +#include "utility.h" + + +namespace zvec { + + +static std::string data_generator_bin_; +const std::string collection_name_{"crash_test"}; +const std::string dir_path_{"crash_test_db"}; +const zvec::CollectionOptions options_{false, true}; + + +static std::string LocateDataGenerator() { + namespace fs = std::filesystem; + const std::vector candidates{"./data_generator", + "./bin/data_generator"}; + for (const auto &p : candidates) { + if (fs::exists(p)) { + return fs::canonical(p).string(); + } + } + throw std::runtime_error("data_generator binary not found"); +} + + +void RunGenerator(const std::string &start, const std::string &end, + const std::string &op, const std::string &version) { + pid_t pid = fork(); + ASSERT_GE(pid, 0); + + if (pid == 0) { // Child process + char arg_path[] = "--path"; + char arg_start[] = "--start"; + char arg_end[] = "--end"; + char arg_op[] = "--op"; + char arg_version[] = "--version"; + char *args[] = {const_cast(data_generator_bin_.c_str()), + arg_path, + const_cast(dir_path_.c_str()), + arg_start, + const_cast(start.c_str()), + arg_end, + const_cast(end.c_str()), + arg_op, + const_cast(op.c_str()), + arg_version, + const_cast(version.c_str()), + nullptr}; + execvp(args[0], args); + perror("execvp failed"); + _exit(1); + } + + int status; + waitpid(pid, &status, 0); + ASSERT_TRUE(WIFEXITED(status)) + << "Child process did not exit normally. Terminated by signal?"; + int exit_code = WEXITSTATUS(status); + ASSERT_EQ(exit_code, 0) << "data_generator failed with exit code: " + << exit_code; +} + + +void RunGeneratorAndCrash(const std::string &start, const std::string &end, + const std::string &op, const std::string &version, + int seconds) { + pid_t pid = fork(); + ASSERT_GE(pid, 0); + + if (pid == 0) { // Child process + char arg_path[] = "--path"; + char arg_start[] = "--start"; + char arg_end[] = "--end"; + char arg_op[] = "--op"; + char arg_version[] = "--version"; + char *args[] = {const_cast(data_generator_bin_.c_str()), + arg_path, + const_cast(dir_path_.c_str()), + arg_start, + const_cast(start.c_str()), + arg_end, + const_cast(end.c_str()), + arg_op, + const_cast(op.c_str()), + arg_version, + const_cast(version.c_str()), + nullptr}; + execvp(args[0], args); + perror("execvp failed"); + _exit(1); + } + + std::this_thread::sleep_for(std::chrono::seconds(seconds)); + if (kill(pid, 0) == 0) { + kill(pid, SIGKILL); + } + int status; + waitpid(pid, &status, 0); + ASSERT_TRUE(WIFSIGNALED(status)) + << "Child process was not killed by a signal. It exited normally?"; +} + + +class CrashRecoveryTest : public ::testing::Test { + protected: + void SetUp() override { + system("rm -rf ./crash_test_db"); + ASSERT_NO_THROW(data_generator_bin_ = LocateDataGenerator()); + } + + void TearDown() override { + system("rm -rf ./crash_test_db"); + } +}; + + +TEST_F(CrashRecoveryTest, BasicInsertAndReopen) { + { + auto schema = CreateTestSchema(collection_name_); + auto result = Collection::CreateAndOpen(dir_path_, *schema, options_); + ASSERT_TRUE(result.has_value()); + auto collection = result.value(); + collection.reset(); + } + + RunGenerator("0", "5000", "insert", "0"); + auto result = Collection::Open(dir_path_, options_); + ASSERT_TRUE(result.has_value()); + auto collection = result.value(); + ASSERT_EQ(collection->Stats().value().doc_count, 5000) + << "Document count mismatch"; +} + + +TEST_F(CrashRecoveryTest, CrashRecoveryDuringInsertion) { + { + auto schema = CreateTestSchema(collection_name_); + auto result = Collection::CreateAndOpen(dir_path_, *schema, options_); + ASSERT_TRUE(result.has_value()); + auto collection = result.value(); + collection.reset(); + } + + RunGeneratorAndCrash("0", "10000", "insert", "0", 3); + + auto result = Collection::Open(dir_path_, options_); + ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. " + "Recovery mechanism may be broken."; + auto collection = result.value(); + uint64_t doc_count{collection->Stats().value().doc_count}; + ASSERT_GT(doc_count, 800) + << "Document count is too low after 3s of insertion and recovery"; + + for (uint64_t doc_id = 0; doc_id < doc_count; doc_id++) { + const auto expected_doc = CreateTestDoc(doc_id, 0); + std::vector pks{}; + pks.emplace_back(expected_doc.pk()); + if (auto res = collection->Fetch(pks); res) { + auto map = res.value(); + if (map.find(expected_doc.pk()) == map.end()) { + FAIL() << "Returned map does not contain doc[" << expected_doc.pk() + << "]"; + } + const auto actual_doc = map.at(expected_doc.pk()); + ASSERT_EQ(*actual_doc, expected_doc) + << "Data mismatch for doc[" << expected_doc.pk() << "]"; + } else { + FAIL() << "Failed to fetch doc[" << expected_doc.pk() << "]"; + } + } +} + + +TEST_F(CrashRecoveryTest, CrashRecoveryDuringUpsert) { + { + auto schema = CreateTestSchema(collection_name_); + auto result = Collection::CreateAndOpen(dir_path_, *schema, options_); + ASSERT_TRUE(result.has_value()); + auto collection = result.value(); + collection.reset(); + } + + RunGenerator("0", "5000", "insert", "0"); + { + auto result = Collection::Open(dir_path_, options_); + ASSERT_TRUE(result.has_value()); + auto collection = result.value(); + ASSERT_EQ(collection->Stats().value().doc_count, 5000) + << "Document count mismatch"; + } + + RunGeneratorAndCrash("4500", "20000", "upsert", "1", 5); + + auto result = Collection::Open(dir_path_, options_); + ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. " + "Recovery mechanism may be broken."; + auto collection = result.value(); + uint64_t doc_count{collection->Stats().value().doc_count}; + ASSERT_GT(doc_count, 6000) + << "Document count is too low after 5s of insertion and recovery"; + + for (uint64_t doc_id = 0; doc_id < doc_count; doc_id++) { + Doc expected_doc; + if (doc_id < 4500) { + expected_doc = CreateTestDoc(doc_id, 0); + } else { + expected_doc = CreateTestDoc(doc_id, 1); + } + std::vector pks{}; + pks.emplace_back(expected_doc.pk()); + if (auto res = collection->Fetch(pks); res) { + auto map = res.value(); + if (map.find(expected_doc.pk()) == map.end()) { + FAIL() << "Returned map does not contain doc[" << expected_doc.pk() + << "]"; + } + const auto actual_doc = map.at(expected_doc.pk()); + ASSERT_EQ(*actual_doc, expected_doc) + << "Data mismatch for doc[" << expected_doc.pk() << "]"; + } else { + FAIL() << "Failed to fetch doc[" << expected_doc.pk() << "]"; + } + } +} + + +TEST_F(CrashRecoveryTest, CrashRecoveryDuringUpdate) { + { + auto schema = CreateTestSchema(collection_name_); + auto result = Collection::CreateAndOpen(dir_path_, *schema, options_); + ASSERT_TRUE(result.has_value()); + auto collection = result.value(); + collection.reset(); + } + + RunGenerator("0", "18000", "upsert", "0"); + { + auto result = Collection::Open(dir_path_, options_); + ASSERT_TRUE(result.has_value()); + auto collection = result.value(); + ASSERT_EQ(collection->Stats().value().doc_count, 18000) + << "Document count mismatch"; + } + + RunGeneratorAndCrash("3000", "15000", "update", "3", 4); + + auto result = Collection::Open(dir_path_, options_); + ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. " + "Recovery mechanism may be broken."; + auto collection = result.value(); + uint64_t doc_count{collection->Stats().value().doc_count}; + ASSERT_EQ(doc_count, 18000) << "Document count mismatch after crash recovery"; + + for (int doc_id = 0; doc_id < 3500; doc_id++) { + Doc expected_doc; + if (doc_id < 3000) { + expected_doc = CreateTestDoc(doc_id, 0); + } else { + expected_doc = CreateTestDoc(doc_id, 3); + } + std::vector pks{}; + pks.emplace_back(expected_doc.pk()); + if (auto res = collection->Fetch(pks); res) { + auto map = res.value(); + if (map.find(expected_doc.pk()) == map.end()) { + FAIL() << "Returned map does not contain doc[" << expected_doc.pk() + << "]"; + } + const auto actual_doc = map.at(expected_doc.pk()); + ASSERT_EQ(*actual_doc, expected_doc) + << "Data mismatch for doc[" << expected_doc.pk() << "]"; + } else { + FAIL() << "Failed to fetch doc[" << expected_doc.pk() << "]"; + } + } +} + + +TEST_F(CrashRecoveryTest, CrashRecoveryDuringDelete) { + { + auto schema = CreateTestSchema(collection_name_); + auto result = Collection::CreateAndOpen(dir_path_, *schema, options_); + ASSERT_TRUE(result.has_value()); + auto collection = result.value(); + collection.reset(); + } + + RunGenerator("0", "18000", "insert", "0"); + { + auto result = Collection::Open(dir_path_, options_); + ASSERT_TRUE(result.has_value()); + auto collection = result.value(); + ASSERT_EQ(collection->Stats().value().doc_count, 18000) + << "Document count mismatch"; + } + + RunGeneratorAndCrash("3000", "15000", "delete", "0", 4); + + auto result = Collection::Open(dir_path_, options_); + ASSERT_TRUE(result.has_value()) << "Failed to reopen collection after crash. " + "Recovery mechanism may be broken."; + auto collection = result.value(); + uint64_t doc_count{collection->Stats().value().doc_count}; + ASSERT_LT(doc_count, 18000) + << "No deletes appear to have been applied before the crash"; + ASSERT_GT(doc_count, 6000) + << "Too many documents deleted, recovery likely lost data"; + + for (int doc_id = 0; doc_id < 3500; doc_id++) { + auto expected_doc = CreateTestDoc(doc_id, 0); + std::vector pks{}; + pks.emplace_back(expected_doc.pk()); + if (auto res = collection->Fetch(pks); res) { + auto map = res.value(); + auto it = map.find(expected_doc.pk()); + ASSERT_NE(it, map.end()) + << "Fetch result missing requested pk[" << expected_doc.pk() << "]"; + if (doc_id < 3000) { + ASSERT_NE(it->second, nullptr) + << "Existing doc returned as nullptr [" << expected_doc.pk() << "]"; + const auto actual_doc = map.at(expected_doc.pk()); + ASSERT_EQ(*actual_doc, expected_doc) + << "Data mismatch for doc[" << expected_doc.pk() << "]"; + } else { + ASSERT_EQ(it->second, nullptr) + << "Returned doc for deleted pk[" << expected_doc.pk() << "]"; + } + } else { + FAIL() << "Failed to fetch doc[" << expected_doc.pk() << "]"; + } + } +} + + +} // namespace zvec