fix: clean up crash residue (#208)
This commit is contained in:
Qinren Zhou
GitHub
parent
ba86c1663d
commit
331933c75a
|
|
@ -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 <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
|
|
@ -1108,6 +1109,52 @@ std::string Doc::to_detail_string() const {
|
|||
return oss.str();
|
||||
}
|
||||
|
||||
struct Doc::ValueEqual {
|
||||
template <typename T, typename U>
|
||||
bool operator()(const T &, const U &) const {
|
||||
return false;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
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<float16_t> &a,
|
||||
const std::vector<float16_t> &b) const {
|
||||
if (a.size() != b.size()) return false;
|
||||
for (size_t i = 0; i < a.size(); ++i)
|
||||
if (std::fabs(static_cast<float>(a[i]) - static_cast<float>(b[i])) >=
|
||||
1e-3f)
|
||||
return false;
|
||||
return true;
|
||||
}
|
||||
|
||||
bool operator()(const std::vector<float> &a,
|
||||
const std::vector<float> &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<double> &a,
|
||||
const std::vector<double> &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<decltype(lhs)>,
|
||||
std::decay_t<decltype(rhs)>>) {
|
||||
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;
|
||||
|
|
|
|||
|
|
@ -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<VectorColumnIndexer>(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<MemForwardStore>(
|
||||
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());
|
||||
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -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 <unistd.h>
|
||||
#include <filesystem>
|
||||
#include <thread>
|
||||
#include <zvec/db/collection.h>
|
||||
#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 <collection_path> --start <start_id> --end <end_id> "
|
||||
"--op <operation>"
|
||||
<< 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<zvec::Doc> 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<zvec::WriteResults> 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<std::string> 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;
|
||||
}
|
||||
|
|
@ -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 <zvec/db/collection.h>
|
||||
#include <zvec/db/doc.h>
|
||||
|
||||
|
||||
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<CollectionSchema>(name);
|
||||
schema->set_max_doc_count_per_segment(2000);
|
||||
|
||||
schema->add_field(
|
||||
std::make_shared<FieldSchema>("int32_field", DataType::INT32, false));
|
||||
schema->add_field(
|
||||
std::make_shared<FieldSchema>("int64_field", DataType::INT64, true));
|
||||
schema->add_field(
|
||||
std::make_shared<FieldSchema>("float_field", DataType::FLOAT, true));
|
||||
schema->add_field(
|
||||
std::make_shared<FieldSchema>("string_field", DataType::STRING, false));
|
||||
schema->add_field(
|
||||
std::make_shared<FieldSchema>("bool_field", DataType::BOOL, false));
|
||||
schema->add_field(std::make_shared<FieldSchema>("array_int32_field",
|
||||
DataType::ARRAY_INT32, true));
|
||||
schema->add_field(std::make_shared<FieldSchema>(
|
||||
"array_string_field", DataType::ARRAY_STRING, false));
|
||||
schema->add_field(std::make_shared<FieldSchema>(
|
||||
"dense_fp32_field", DataType::VECTOR_FP32, 128, false,
|
||||
std::make_shared<HnswIndexParams>(MetricType::COSINE)));
|
||||
schema->add_field(std::make_shared<FieldSchema>(
|
||||
"sparse_fp32_field", DataType::SPARSE_VECTOR_FP32, 0, false,
|
||||
std::make_shared<HnswIndexParams>(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_t>("int32_field", static_cast<int32_t>(doc_id));
|
||||
|
||||
// int64_field: nullable, null if doc_id % 60 == 0
|
||||
if (doc_id % 60 != 0) {
|
||||
doc.set<int64_t>("int64_field", static_cast<int64_t>(doc_id));
|
||||
}
|
||||
|
||||
// float_field: nullable, null if doc_id % 70 == 0
|
||||
if (doc_id % 70 != 0) {
|
||||
doc.set<float>("float_field", static_cast<float>(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<std::string>("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>("bool_field", bool_value);
|
||||
|
||||
// array_int32_field: nullable, null if doc_id % 100 == 0
|
||||
if (doc_id % 100 != 0) {
|
||||
std::vector<int32_t> array_int32;
|
||||
for (int i = 0; i < 3; i++) {
|
||||
array_int32.push_back(static_cast<int32_t>(doc_id + i));
|
||||
}
|
||||
doc.set<std::vector<int32_t>>("array_int32_field", array_int32);
|
||||
}
|
||||
|
||||
// array_string_field: ["str_0", "str_1", ...] or ["updated_str_0", ...]
|
||||
std::vector<std::string> 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<std::vector<std::string>>("array_string_field", array_string);
|
||||
|
||||
// dense_fp32_field: deterministic pattern
|
||||
std::vector<float> dense(128);
|
||||
for (int i = 0; i < 128; i++) {
|
||||
dense[i] = static_cast<float>(doc_id + i) / 1000.0f;
|
||||
}
|
||||
doc.set<std::vector<float>>("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<uint32_t> sparse_indices;
|
||||
std::vector<float> sparse_values;
|
||||
for (uint32_t idx = 0; idx <= 100; idx += 10) {
|
||||
sparse_indices.push_back(idx);
|
||||
sparse_values.push_back(static_cast<float>(doc_id + idx) / 1000.0f);
|
||||
}
|
||||
doc.set<std::pair<std::vector<uint32_t>, std::vector<float>>>(
|
||||
"sparse_fp32_field", std::make_pair(sparse_indices, sparse_values));
|
||||
|
||||
return doc;
|
||||
}
|
||||
|
||||
|
||||
} // namespace zvec
|
||||
|
|
@ -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 <csignal>
|
||||
#include <filesystem>
|
||||
#include <thread>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/db/collection.h>
|
||||
#include <zvec/db/doc.h>
|
||||
#include <zvec/db/schema.h>
|
||||
#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<std::string> 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<char *>(data_generator_bin_.c_str()),
|
||||
arg_path,
|
||||
const_cast<char *>(dir_path_.c_str()),
|
||||
arg_start,
|
||||
const_cast<char *>(start.c_str()),
|
||||
arg_end,
|
||||
const_cast<char *>(end.c_str()),
|
||||
arg_op,
|
||||
const_cast<char *>(op.c_str()),
|
||||
arg_version,
|
||||
const_cast<char *>(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<char *>(data_generator_bin_.c_str()),
|
||||
arg_path,
|
||||
const_cast<char *>(dir_path_.c_str()),
|
||||
arg_start,
|
||||
const_cast<char *>(start.c_str()),
|
||||
arg_end,
|
||||
const_cast<char *>(end.c_str()),
|
||||
arg_op,
|
||||
const_cast<char *>(op.c_str()),
|
||||
arg_version,
|
||||
const_cast<char *>(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<std::string> 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<std::string> 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<std::string> 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<std::string> 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
|
||||
Loading…
Reference in New Issue