From e2130ebbb63decd1f96535fe9e4cd7a853830098 Mon Sep 17 00:00:00 2001 From: Qinren Zhou Date: Wed, 16 Sep 2026 15:08:23 +0800 Subject: [PATCH 1/2] fix: handle WAL corruption and recover upserts correctly --- src/db/index/segment/segment.cc | 299 +++++++--- src/db/index/storage/wal/local_wal_file.cc | 240 ++++---- src/db/index/storage/wal/local_wal_file.h | 15 +- src/db/index/storage/wal/wal_file.h | 9 +- tests/db/crash_recovery/wal_recovery_test.cc | 562 +++++++++++++++++++ tests/db/index/storage/wal_file_test.cc | 231 ++++++-- 6 files changed, 1090 insertions(+), 266 deletions(-) create mode 100644 tests/db/crash_recovery/wal_recovery_test.cc diff --git a/src/db/index/segment/segment.cc b/src/db/index/segment/segment.cc index efc27f758..26742c8d7 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -317,10 +317,12 @@ class SegmentImpl : public Segment, Status insert_vector_indexer(Doc &doc); Status internal_insert(Doc &doc); Status internal_update(Doc &doc); - Status internal_upsert(Doc &doc); Status internal_delete(const Doc &doc); Status recover(); + Result> find_legacy_upsert_predecessors( + const std::unordered_set &keys, + uint64_t first_replay_id) const; Status open_wal_file(); Status append_wal(const Doc &doc); Status update_version(uint32_t delete_snapshot_path_suffix); @@ -950,15 +952,6 @@ Status SegmentImpl::internal_update(Doc &doc) { return internal_insert(doc); } -Status SegmentImpl::internal_upsert(Doc &doc) { - uint64_t g_doc_id; - bool exist = id_map_->has(doc.pk_ref(), &g_doc_id); - if (exist) { - delete_store_->mark_deleted(g_doc_id); - } - return internal_insert(doc); -} - Status SegmentImpl::internal_delete(const Doc &doc) { delete_store_->mark_deleted(doc.doc_id()); id_map_->remove(doc.pk_ref()); @@ -1003,13 +996,28 @@ Status SegmentImpl::Update(Doc &doc) { Status SegmentImpl::Upsert(Doc &doc) { std::lock_guard lock(seg_mtx_); - doc.set_operator(Operator::UPSERT); - - // append WAL - auto s = append_wal(doc); - CHECK_RETURN_STATUS(s); + // Persist the predecessor explicitly. RocksDB may flush the new ID mapping + // before the deletion snapshot is committed, so recovery cannot safely + // infer the superseded document from the current ID map. + const auto original_doc_id = doc.doc_id(); + uint64_t previous_id; + const bool exists = id_map_->has(doc.pk_ref(), &previous_id); + if (exists) { + doc.set_doc_id(previous_id); + doc.set_operator(Operator::UPDATE); + } else { + doc.set_operator(Operator::INSERT); + } - return internal_upsert(doc); + auto status = append_wal(doc); + // Preserve the public operation on the caller's document; only the WAL + // uses the already-supported INSERT/UPDATE representation. + doc.set_operator(Operator::UPSERT); + if (!status.ok()) { + doc.set_doc_id(original_doc_id); + return status; + } + return exists ? internal_update(doc) : internal_insert(doc); } Status SegmentImpl::Delete(const std::string &pk) { @@ -4266,10 +4274,82 @@ Status SegmentImpl::init_memory_components() { return Status::OK(); } +Result> SegmentImpl::find_legacy_upsert_predecessors( + const std::unordered_set &keys, + uint64_t first_replay_id) const { + std::vector predecessors; + const auto version = version_manager_->get_current_version(); + auto segments = version.persisted_segment_metas(); + if (auto writing = version.writing_segment_meta()) { + segments.push_back(std::move(writing)); + } + for (const auto &segment : segments) { + for (const auto &block : segment->persisted_blocks()) { + if (block.type() != BlockType::SCALAR || + !block.contain_column(GLOBAL_DOC_ID) || + !block.contain_column(USER_ID)) { + continue; + } + const auto forward_path = FileHelper::MakeForwardBlockPath( + path_, segment->id(), block.id(), !options_.enable_mmap_); + BaseForwardStore::Ptr store; + // BufferPoolForwardStore only supports Parquet; IPC always uses the + // mapped store. Release each store and its buffers before the next block. + if (options_.enable_mmap_ || + InferFileFormat(forward_path) == FileFormat::IPC) { + store = std::make_shared(forward_path); + } else { + store = std::make_shared(forward_path); + } + auto status = store->Open(); + if (!status.ok()) { + return tl::make_unexpected(Status::InternalError( + "Failed to open committed rows for legacy WAL recovery: path[", + forward_path, "], reason[", status.message(), "]")); + } + auto reader = store->scan({GLOBAL_DOC_ID, USER_ID}); + if (!reader) { + return tl::make_unexpected(Status::InternalError( + "Failed to scan committed rows for legacy WAL recovery: ", + forward_path)); + } + while (true) { + std::shared_ptr batch; + const auto read_status = reader->ReadNext(&batch); + if (!read_status.ok()) { + return tl::make_unexpected(Status::InternalError( + "Failed to read committed rows for legacy WAL recovery: path[", + forward_path, "], reason[", read_status.message(), "]")); + } + if (!batch) break; + const auto ids = std::dynamic_pointer_cast( + batch->GetColumnByName(GLOBAL_DOC_ID)); + const auto pks = std::dynamic_pointer_cast( + batch->GetColumnByName(USER_ID)); + if (!ids || !pks || ids->length() != pks->length() || + ids->null_count() != 0 || pks->null_count() != 0) { + return tl::make_unexpected(Status::InternalError( + "Invalid committed identity columns during legacy WAL recovery: ", + forward_path)); + } + for (int64_t row = 0; row < ids->length(); ++row) { + const auto doc_id = ids->Value(row); + if (doc_id < first_replay_id && !delete_store_->is_deleted(doc_id) && + keys.find(pks->GetString(row)) != keys.end()) { + predecessors.push_back(doc_id); + } + } + } + } + } + return predecessors; +} + Status SegmentImpl::recover() { // recover mem block meta auto &mem_block = segment_meta_->writing_forward_block().value(); - doc_id_allocator_.store(mem_block.min_doc_id()); + const auto first_replay_id = mem_block.min_doc_id(); + doc_id_allocator_.store(first_replay_id); std::string wal_file_path = FileHelper::MakeWalPath(path_, segment_meta_->id(), mem_block.id_); @@ -4286,99 +4366,141 @@ Status SegmentImpl::recover() { 0) { LOG_ERROR("WAL recovery failed: unable to open WAL file [%s]", wal_file_path.c_str()); - return Status::OK(); + return Status::InternalError("Failed to open WAL for recovery: ", + wal_file_path); } std::array(Operator::DELETE) + 1> recovered_doc_count{}; uint64_t total_recovered_doc_count{0}; - - int ret = recover_wal_file->prepare_for_read(); - if (ret != 0) { - LOG_ERROR( - "WAL recovery failed: unable to prepare file for reading, path[%s], " - "segment[%d], ret[%d]", - wal_file_path.c_str(), id(), ret); - return Status::InternalError( - "Failed to prepare WAL file for reading: path[", wal_file_path, - "], segment[", id(), "], ret[", ret, "]"); - } + std::unordered_set legacy_upsert_keys; LOG_INFO("WAL recovery started: path[%s], segment[%d]", wal_file_path.c_str(), id()); std::lock_guard lock(seg_mtx_); - while (true) { - std::string buf = recover_wal_file->next(); - if (buf.empty()) { - break; - } - total_recovered_doc_count++; - auto doc = Doc::deserialize(reinterpret_cast(buf.data()), - buf.size()); - if (doc == nullptr) { + // Validate the complete stream before changing the ID map or indexes. A + // failed open can close and flush those stores, so discovering corruption + // after applying a prefix would otherwise persist a partial recovery. + // Read one WAL record at a time. Legacy recovery additionally holds unique + // UPSERT keys, matched predecessor IDs, and the existing forward-store + // buffers for one committed block. New WAL records need no committed scan. + for (int pass = 0; pass < 2; ++pass) { + const bool replay = pass == 1; + total_recovered_doc_count = 0; + int ret = recover_wal_file->prepare_for_read(); + if (ret != 0) { LOG_ERROR( - "WAL record recovery failed: path[%s], segment[%d], record[%zu], " - "reason[deserialization failed]", - wal_file_path.c_str(), id(), (size_t)total_recovered_doc_count); - continue; + "WAL recovery failed: unable to prepare file for reading, path[%s], " + "segment[%d], ret[%d]", + wal_file_path.c_str(), id(), ret); + return Status::InternalError( + "Failed to prepare WAL file for reading: path[", wal_file_path, + "], segment[", id(), "], ret[", ret, "]"); + } + + if (replay && !legacy_upsert_keys.empty()) { + // Old UPSERT records did not include their predecessor ID. Reconstruct + // it from committed rows even if an interrupted replay already replaced + // its ID-map entry. Finish the entire scan before changing tombstones. + auto predecessors = + find_legacy_upsert_predecessors(legacy_upsert_keys, first_replay_id); + if (!predecessors.has_value()) return predecessors.error(); + for (const auto doc_id : predecessors.value()) { + delete_store_->mark_deleted(doc_id); + } } - Status status; - switch (doc->get_operator()) { - case Operator::INSERT: { - internal_insert(*doc); - break; - } - case Operator::UPDATE: { - internal_update(*doc); - break; + while (true) { + auto record = recover_wal_file->next(); + if (!record.has_value()) { + return Status::InternalError( + "Failed to read WAL during recovery: path[", wal_file_path, + "], segment[", id(), "], reason[", record.error().message(), "]"); } - case Operator::UPSERT: { - internal_upsert(*doc); + if (!record.value().has_value()) { break; } - case Operator::DELETE: { - internal_delete(*doc); - break; + if (options_.read_only_) { + return Status::FailedPrecondition( + "WAL recovery is required; open the collection in read-write mode " + "once to recover before opening it read-only"); } - default: + const auto &buf = record.value().value(); + total_recovered_doc_count++; + auto doc = Doc::deserialize(reinterpret_cast(buf.data()), + buf.size()); + if (doc == nullptr) { LOG_ERROR( "WAL record recovery failed: path[%s], segment[%d], record[%zu], " - "operator[%d], reason[unknown operator]", - wal_file_path.c_str(), id(), (size_t)total_recovered_doc_count, - static_cast(doc->get_operator())); - break; - } + "reason[deserialization failed]", + wal_file_path.c_str(), id(), (size_t)total_recovered_doc_count); + return Status::InternalError( + "Corrupt WAL document: path[", wal_file_path, "], segment[", id(), + "], record[", total_recovered_doc_count, "]"); + } - if (!status.ok()) { - LOG_ERROR( - "WAL record recovery failed: path[%s], segment[%d], record[%zu], " - "operator[%d], reason[%s]", - wal_file_path.c_str(), id(), (size_t)total_recovered_doc_count, - static_cast(doc->get_operator()), status.message().c_str()); - continue; - } + if (!replay) { + if (doc->get_operator() == Operator::UPSERT) { + legacy_upsert_keys.insert(doc->pk_ref()); + } + continue; + } - recovered_doc_count[static_cast(doc->get_operator())]++; - } + Status status; + switch (doc->get_operator()) { + case Operator::INSERT: { + status = internal_insert(*doc); + break; + } + case Operator::UPDATE: { + status = internal_update(*doc); + break; + } + case Operator::UPSERT: { + // A previous interrupted replay may already have persisted this + // record's ID (or a later ID for the same key) in RocksDB. Only an + // older document is superseded; marking this/later replay ID deleted + // would hide a successfully recovered document on retry. + uint64_t previous_id; + if (id_map_->has(doc->pk_ref(), &previous_id) && + previous_id < doc_id_allocator_.load()) { + delete_store_->mark_deleted(previous_id); + } + status = internal_insert(*doc); + break; + } + case Operator::DELETE: { + status = internal_delete(*doc); + break; + } + default: + LOG_ERROR( + "WAL record recovery failed: path[%s], segment[%d], record[%zu], " + "operator[%d], reason[unknown operator]", + wal_file_path.c_str(), id(), (size_t)total_recovered_doc_count, + static_cast(doc->get_operator())); + return Status::InternalError("Unknown WAL document operator: path[", + wal_file_path, "], record[", + total_recovered_doc_count, "]"); + } - const auto added_docs = recovered_doc_count[0] + // INSERT - recovered_doc_count[1] + // UPSERT - recovered_doc_count[2]; // UPDATE - mem_block.max_doc_id_ += added_docs; + if (!status.ok()) { + LOG_ERROR( + "WAL record recovery failed: path[%s], segment[%d], record[%zu], " + "operator[%d], reason[%s]", + wal_file_path.c_str(), id(), (size_t)total_recovered_doc_count, + static_cast(doc->get_operator()), status.message().c_str()); + return Status(status.code(), + ailego::StringHelper::Concat( + "Failed to apply WAL record: path[", wal_file_path, + "], record[", total_recovered_doc_count, "], reason[", + status.message(), "]")); + } - ret = recover_wal_file->close(); - if (ret != 0) { - LOG_ERROR( - "WAL recovery failed: unable to close file, path[%s], " - "segment[%d], ret[%d]", - wal_file_path.c_str(), id(), ret); - return Status::InternalError("Failed to close recovered WAL file: path[", - wal_file_path, "], segment[", id(), "], ret[", - ret, "]"); + recovered_doc_count[static_cast(doc->get_operator())]++; + } } - recover_wal_file.reset(); LOG_INFO( "WAL recovery completed: path[%s], segment[%d], total[%zu], " @@ -4394,7 +4516,10 @@ Status SegmentImpl::recover() { // optimize() flush the writing segment before sealing it; without an open // member WAL, flush() treats the recovered memory components as empty and // returns without persisting them. - return open_wal_file(); + // Retain the reader's valid-tail position so a later append can discard an + // incomplete crash record. Opening read-only does not truncate the WAL. + wal_file_ = std::move(recover_wal_file); + return Status::OK(); } Status SegmentImpl::open_wal_file() { diff --git a/src/db/index/storage/wal/local_wal_file.cc b/src/db/index/storage/wal/local_wal_file.cc index 6a50e5c34..e9d0493c5 100644 --- a/src/db/index/storage/wal/local_wal_file.cc +++ b/src/db/index/storage/wal/local_wal_file.cc @@ -13,6 +13,9 @@ // limitations under the License. #include "local_wal_file.h" +#include +#include +#include #ifndef _MSC_VER #include #endif @@ -22,49 +25,73 @@ #include "db/common/file_helper.h" #include "db/common/typedef.h" -#define MAX_RECORD_SIZE 4194304 // 4Mb - namespace zvec { int LocalWalFile::append(std::string &&data) { + if (data.empty() || data.size() > std::numeric_limits::max()) { + WLOG_ERROR("Wal record length is not representable: %zu", data.size()); + return -1; + } + WalRecord record; - record.length_ = data.size(); - record.crc_ = ailego::Crc32c::Hash( - reinterpret_cast(data.data()), record.length_, 0); - record.content_ = std::forward(data); + record.length_ = static_cast(data.size()); + record.crc_ = ailego::Crc32c::Hash(data.data(), data.size(), 0); + record.content_ = std::move(data); + std::lock_guard lock(file_mutex_); + if (!opened_ || failed_) { + return -1; + } + if (incomplete_tail_offset_) { + if (!file_.truncate(*incomplete_tail_offset_)) { + WLOG_ERROR("Wal incomplete tail truncation failed"); + failed_ = true; + return -1; + } + incomplete_tail_offset_.reset(); + } + if (!file_.seek(0, ailego::File::Origin::End)) { + return -1; + } if (write_record(record) < 0) { - WLOG_ERROR("Wal write record error. record.length_[%zu]", - (size_t)record.length_); return -1; } - // if max_docs_wal_flush_ is 0, no need flush + // Keep the flush counter and flush in the same critical section as writes. if (max_docs_wal_flush_ != 0 && docs_count_ >= max_docs_wal_flush_) { if (!file_.flush()) { WLOG_ERROR("Wal flush error. docs_count_[%zu] max_docs_wal_flush_[%zu]", (size_t)docs_count_, (size_t)max_docs_wal_flush_); + failed_ = true; + return -1; } docs_count_ = 0; } return 0; } -std::string LocalWalFile::next() { +Result> LocalWalFile::next() { + std::lock_guard lock(file_mutex_); + if (!opened_ || failed_) { + return tl::make_unexpected( + Status::InternalError("WAL is not open for reading or has failed")); + } WalRecord record; - if (read_record(record) > 0) { - uint32_t tmp_crc = ailego::Crc32c::Hash( - reinterpret_cast(record.content_.data()), record.length_, - 0); - if (tmp_crc == record.crc_) { - return std::move(record.content_); - } else { - WLOG_ERROR( - "Wal next error. record.length_[%zu] crc_[%zu] != tmp_crc[%zu]", - (size_t)record.length_, (size_t)record.crc_, (size_t)tmp_crc); - } + auto result = read_record(record); + if (!result.has_value()) { + failed_ = true; + return tl::make_unexpected(result.error()); + } + if (!result.value()) { + return std::nullopt; + } + const uint32_t crc = + ailego::Crc32c::Hash(record.content_.data(), record.content_.size(), 0); + if (crc != record.crc_) { + failed_ = true; + return tl::make_unexpected( + Status::InternalError("WAL record CRC mismatch")); } - // end of file or read error - return std::string(); + return std::optional(std::move(record.content_)); } int LocalWalFile::open(const WalOptions &wal_option) { @@ -82,7 +109,7 @@ int LocalWalFile::open(const WalOptions &wal_option) { } // write wal header - int write_size = file_.write((const void *)&header_, sizeof(header_)); + size_t write_size = file_.write((const void *)&header_, sizeof(header_)); if (write_size != sizeof(header_)) { WLOG_ERROR("Wal write header error. create_new[%d]", wal_option.create_new); @@ -102,11 +129,16 @@ int LocalWalFile::open(const WalOptions &wal_option) { } // open default for write - file_.seek(0, ailego::File::Origin::End); + if (!file_.seek(0, ailego::File::Origin::End)) { + return -1; + } } max_docs_wal_flush_ = wal_option.max_docs_wal_flush; opened_ = true; + failed_ = false; + incomplete_tail_offset_.reset(); + docs_count_ = 0; WLOG_INFO("Wal open success. create_new[%d]", wal_option.create_new); return 0; @@ -142,106 +174,98 @@ int LocalWalFile::flush() { int LocalWalFile::prepare_for_read() { CHECK_STATUS(opened_, true); - if (!file_.seek(0, ailego::File::Origin::Begin)) { + incomplete_tail_offset_.reset(); + if (failed_ || !file_.seek(0, ailego::File::Origin::Begin)) { return -1; } - int read_size = file_.read((void *)&header_, sizeof(header_)); + size_t read_size = file_.read((void *)&header_, sizeof(header_)); if (read_size != sizeof(header_)) { WLOG_ERROR("Wal read header error."); + failed_ = true; return -1; } if (header_.wal_version != 0UL) { WLOG_ERROR("Wal version not support error."); + failed_ = true; return -1; } return 0; } -//! Return 1 if success or -1 if write error +// Caller holds file_mutex_. A failed write must not strand future successful +// appends behind its incomplete record. int LocalWalFile::write_record(WalRecord &record) { - CHECK_STATUS(opened_, true); - - int write_size = 0; - int ret = -1; - - std::lock_guard lock(file_mutex_); - do { - write_size = file_.write((const void *)&record.length_, LENGTH_SIZE); - if (write_size != LENGTH_SIZE) { - WLOG_ERROR("Wal write error. record.length_ error write_size[%d]", - write_size); - break; - } - - write_size = file_.write((const void *)&record.crc_, CRC_SIZE); - if (write_size != CRC_SIZE) { - WLOG_ERROR("Wal write error. record.crc_ error write_size[%d]", - write_size); - break; - } - - write_size = - file_.write((const void *)record.content_.data(), record.length_); - if (write_size != (int)record.length_) { - WLOG_ERROR("Wal write error. record.content_ error write_size[%d]", - write_size); - break; + const auto start = file_.offset(); + if (start < static_cast(sizeof(header_))) { + failed_ = true; + return -1; + } + if (file_.write(&record.length_, LENGTH_SIZE) != LENGTH_SIZE || + file_.write(&record.crc_, CRC_SIZE) != CRC_SIZE || + file_.write(record.content_.data(), record.content_.size()) != + record.content_.size()) { + WLOG_ERROR("Wal write record failed. record.length_[%zu]", + record.content_.size()); + if (!file_.truncate(static_cast(start)) || + !file_.seek(start, ailego::File::Origin::Begin)) { + failed_ = true; } - ret = 1; // write one record success - docs_count_++; - } while (false); - - return ret; + return -1; + } + ++docs_count_; + return 1; } -//! Return 1 if success or 0 if eof or -1 if read error -int LocalWalFile::read_record(WalRecord &record) { - CHECK_STATUS(opened_, true); - - int read_size = 0; - std::string err_msg; - int ret = -1; - - do { - read_size = - file_.read(reinterpret_cast(&record.length_), LENGTH_SIZE); - if (read_size == 0) { - ret = 0; - WLOG_INFO("Wal read finished. end of file"); - break; - } - - if (read_size != LENGTH_SIZE) { - WLOG_ERROR("Wal read error. record.length_ error read_size[%d]", - read_size); - break; - } - - read_size = file_.read(reinterpret_cast(&record.crc_), CRC_SIZE); - if (read_size != CRC_SIZE) { - WLOG_ERROR("Wal read error. record.crc_ error read_size[%d]", read_size); - break; - } - - // resize may crash if record.length_ very large - if (record.length_ <= 0 || record.length_ > MAX_RECORD_SIZE) { - WLOG_ERROR("Wal read error. record.length_ value error read_size[%d]", - read_size); - break; - } - +Result LocalWalFile::read_record(WalRecord &record) { + if (incomplete_tail_offset_) { + return false; + } + // File::read reports bytes read for both EOF and I/O failures. Check the + // physical extent first: a short read within that extent is an I/O error, + // whereas a final frame that does not fit is a tolerated interrupted write. + const auto start = file_.offset(); + const size_t file_size = file_.size(); + if (!file_.is_valid() || start < static_cast(sizeof(header_)) || + file_size < sizeof(header_) || static_cast(start) > file_size) { + return tl::make_unexpected( + Status::InternalError("Failed to determine WAL read position or size")); + } + const size_t remaining = file_size - static_cast(start); + if (remaining == 0) { + return false; + } + if (remaining < LENGTH_SIZE + CRC_SIZE) { + incomplete_tail_offset_ = static_cast(start); + return false; + } + if (file_.read(&record.length_, LENGTH_SIZE) != LENGTH_SIZE || + file_.read(&record.crc_, CRC_SIZE) != CRC_SIZE) { + return tl::make_unexpected( + Status::InternalError("Failed to read WAL record header")); + } + if (record.length_ == 0) { + return tl::make_unexpected( + Status::InternalError("WAL record has zero length")); + } + if (record.length_ > remaining - LENGTH_SIZE - CRC_SIZE) { + incomplete_tail_offset_ = static_cast(start); + return false; + } + try { record.content_.resize(record.length_); - read_size = file_.read((void *)const_cast(record.content_.data()), - record.length_); - if (read_size != (int)record.length_) { - WLOG_ERROR("Wal read error. record.content_ error read_size[%d]", - read_size); - break; - } - ret = 1; // read one record success - } while (false); - - return ret; + } catch (const std::bad_alloc &) { + return tl::make_unexpected(Status(StatusCode::RESOURCE_EXHAUSTED, + "Unable to allocate WAL record buffer")); + } catch (const std::length_error &) { + return tl::make_unexpected(Status(StatusCode::RESOURCE_EXHAUSTED, + "WAL record exceeds string capacity")); + } + if (file_.read(record.content_.data(), record.content_.size()) != + record.content_.size()) { + return tl::make_unexpected( + Status::InternalError("Failed to read WAL record payload")); + } + return true; } -}; // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/index/storage/wal/local_wal_file.h b/src/db/index/storage/wal/local_wal_file.h index d4f392fa8..c2a8dcadf 100644 --- a/src/db/index/storage/wal/local_wal_file.h +++ b/src/db/index/storage/wal/local_wal_file.h @@ -14,12 +14,7 @@ #pragma once #include -#include -#include -#include #include -#include -#include #include #include "wal_file.h" @@ -30,7 +25,7 @@ namespace zvec { */ struct WalHeader { uint64_t wal_version{0U}; - uint64_t reserved_[7]; + uint64_t reserved_[7]{}; }; static_assert(sizeof(WalHeader) % 64 == 0, @@ -61,7 +56,7 @@ class LocalWalFile : public WalFile { public: int append(std::string &&data) override; int prepare_for_read() override; - std::string next() override; + Result> next() override; public: int open(const WalOptions &wal_option) override; @@ -78,7 +73,7 @@ class LocalWalFile : public WalFile { private: int write_record(WalRecord &record); - int read_record(WalRecord &record); + Result read_record(WalRecord &record); private: ailego::File file_; @@ -93,6 +88,10 @@ class LocalWalFile : public WalFile { WalHeader header_; bool opened_{false}; + bool failed_{false}; + // Preserve the complete prefix and remove a torn final record before the + // next append. Merely reading a WAL must not modify it. + std::optional incomplete_tail_offset_; }; diff --git a/src/db/index/storage/wal/wal_file.h b/src/db/index/storage/wal/wal_file.h index 6e17f7384..34c2570b3 100644 --- a/src/db/index/storage/wal/wal_file.h +++ b/src/db/index/storage/wal/wal_file.h @@ -13,8 +13,11 @@ // limitations under the License. #pragma once +#include #include +#include #include +#include namespace zvec { @@ -46,7 +49,9 @@ class WalFile { public: virtual int append(std::string &&data) = 0; virtual int prepare_for_read() = 0; - virtual std::string next() = 0; + // A successful empty optional means EOF or an incomplete final crash record. + // Read failures and complete but corrupt records return an error. + virtual Result> next() = 0; public: //! Open and initialize WalFile @@ -64,4 +69,4 @@ class WalFile { virtual bool has_record() = 0; }; -}; // namespace zvec \ No newline at end of file +}; // namespace zvec diff --git a/tests/db/crash_recovery/wal_recovery_test.cc b/tests/db/crash_recovery/wal_recovery_test.cc new file mode 100644 index 000000000..ef757408c --- /dev/null +++ b/tests/db/crash_recovery/wal_recovery_test.cc @@ -0,0 +1,562 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "db/index/common/id_map.h" +#include "db/index/common/version_manager.h" +#include "db/index/storage/wal/wal_file.h" + +namespace zvec { +namespace { + +std::string ReadFileBytes(const std::string &path) { + std::ifstream file(path, std::ios::binary); + return std::string(std::istreambuf_iterator(file), {}); +} + +std::string FindWal(const std::string &path) { + for (const auto &entry : + std::filesystem::recursive_directory_iterator(path)) { + if (entry.path().extension() == ".wal") return entry.path().string(); + } + return {}; +} + +std::map ReadManifests(const std::string &path) { + std::map result; + for (const auto &entry : + std::filesystem::recursive_directory_iterator(path)) { + if (entry.path().filename().string().rfind("manifest", 0) == 0) { + result.emplace(entry.path().string(), + ReadFileBytes(entry.path().string())); + } + } + return result; +} + +// Only called in ASSERT_EXIT children. Intentionally skip collection cleanup. +void WriteStringDocsAndExit(const std::string &path, + const std::vector &ids, + const std::vector &values) { + CollectionSchema schema("wal_recovery"); + if (!schema + .add_field( + std::make_shared("text", DataType::STRING, false)) + .ok()) { + std::_Exit(1); + } + auto created = Collection::CreateAndOpen(path, schema, CollectionOptions{}); + if (!created.has_value()) std::_Exit(2); + auto collection = std::move(created).value(); + std::vector docs; + for (size_t i = 0; i < ids.size(); ++i) { + Doc doc; + doc.set_pk(ids[i]); + doc.set("text", values[i]); + docs.push_back(std::move(doc)); + } + auto result = collection->insert(docs); + if (!result.has_value()) std::_Exit(3); + for (const auto &status : result.value()) { + if (!status.ok()) std::_Exit(4); + } + std::_Exit(0); +} + +std::string FindIdMap(const std::string &path) { + for (const auto &entry : std::filesystem::directory_iterator(path)) { + if (entry.is_directory() && + entry.path().filename().string().rfind("idmap", 0) == 0) { + return entry.path().string(); + } + } + return {}; +} + +void WriteUpsertsAndExit(const std::string &path, bool persisted_base, + bool append_suffix, bool separate_segment = false) { + CollectionSchema schema("wal_recovery"); + if (!schema + .add_field( + std::make_shared("text", DataType::STRING, false)) + .ok()) { + std::_Exit(1); + } + auto created = Collection::CreateAndOpen(path, schema, CollectionOptions{}); + if (!created.has_value()) std::_Exit(2); + auto collection = std::move(created).value(); + Doc doc; + doc.set_pk("target"); + if (persisted_base) { + doc.set("text", "original"); + std::vector docs{doc}; + auto inserted = collection->insert(docs); + if (!inserted.has_value() || !inserted.value().front().ok() || + !collection->flush().ok()) + std::_Exit(3); + if (separate_segment && !collection->optimize().ok()) std::_Exit(6); + } + for (const auto &value : {"first", "second"}) { + doc.set("text", value); + std::vector docs{doc}; + auto updated = collection->upsert(docs); + if (!updated.has_value() || !updated.value().front().ok()) std::_Exit(4); + } + if (append_suffix) { + doc.set_pk("broken"); + std::vector docs{doc}; + auto inserted = collection->insert(docs); + if (!inserted.has_value() || !inserted.value().front().ok()) std::_Exit(5); + } + std::_Exit(0); +} + +void ReadWalDocuments(const std::string &path, std::vector *docs) { + const auto wal_path = FindWal(path); + ASSERT_FALSE(wal_path.empty()); + auto wal = WalFile::Create(wal_path); + ASSERT_EQ(wal->open(WalOptions{}), 0); + ASSERT_EQ(wal->prepare_for_read(), 0); + while (true) { + auto record = wal->next(); + ASSERT_TRUE(record.has_value()) << record.error().message(); + if (!record.value().has_value()) break; + const auto &bytes = record.value().value(); + auto doc = Doc::deserialize(reinterpret_cast(bytes.data()), + bytes.size()); + ASSERT_NE(doc, nullptr); + docs->push_back(std::move(doc)); + } + ASSERT_EQ(wal->close(), 0); +} + +// Recreate historical UPSERT records through the serializer and WAL writer so +// the framing and checksums remain valid. New public writes use INSERT/UPDATE. +void RewriteWalAsLegacyUpserts(const std::string &path) { + std::vector docs; + ASSERT_NO_FATAL_FAILURE(ReadWalDocuments(path, &docs)); + ASSERT_FALSE(docs.empty()); + auto wal = WalFile::Create(FindWal(path)); + ASSERT_EQ(wal->remove(), 0); + WalOptions options; + options.create_new = true; + ASSERT_EQ(wal->open(options), 0); + for (auto &doc : docs) { + if (doc->pk_ref() == "target") { + doc->set_operator(Operator::UPSERT); + // Legacy UPSERT did not record its predecessor's ID. + doc->set_doc_id(0); + } + auto bytes = doc->serialize(); + ASSERT_EQ(wal->append(std::string(bytes.begin(), bytes.end())), 0); + } + ASSERT_EQ(wal->flush(), 0); + ASSERT_EQ(wal->close(), 0); +} + +void ExpectOnlyTarget(const Collection::Ptr &collection, + const std::string &value) { + auto fetched = collection->fetch({"target"}); + ASSERT_TRUE(fetched.has_value()) << fetched.error().message(); + ASSERT_NE(fetched.value().at("target"), nullptr); + EXPECT_EQ(fetched.value().at("target")->get("text"), value); + + SearchQuery query; + query.topk_ = 10; + query.filter_ = "text != ''"; + auto matches = collection->query(query); + ASSERT_TRUE(matches.has_value()) << matches.error().message(); + ASSERT_EQ(matches.value().size(), 1u); + EXPECT_EQ(matches.value().front()->pk_ref(), "target"); + EXPECT_EQ(matches.value().front()->get("text"), value); + auto stats = collection->stats(); + ASSERT_TRUE(stats.has_value()) << stats.error().message(); + EXPECT_EQ(stats.value().doc_count, 1u); +} + +class WalRecoveryDeathTest : public ::testing::Test { + protected: + void SetUp() override { + // Re-exec children start with the default style; select threadsafe before + // InDeathTestChild() interprets their internal death-test flag. + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + // Threadsafe death tests re-exec the fixture. A second crash child must + // reopen the first child's database instead of deleting it in SetUp. + if (!::testing::internal::InDeathTestChild()) { + ailego::FileHelper::RemovePath(path_.c_str()); + } + } + + void TearDown() override { + ailego::FileHelper::RemovePath(path_.c_str()); + } + + void CheckLegacyCommittedPredecessor(bool separate_segment) { + ASSERT_EXIT(WriteUpsertsAndExit(path_, true, false, separate_segment), + ::testing::ExitedWithCode(0), ""); + if (!::testing::internal::InDeathTestChild()) { + auto recovered_version = VersionManager::Recovery(path_); + ASSERT_TRUE(recovered_version.has_value()); + const auto version = recovered_version.value()->get_current_version(); + EXPECT_EQ(version.persisted_segment_metas().empty(), !separate_segment); + if (!separate_segment) { + EXPECT_FALSE( + version.writing_segment_meta()->persisted_blocks().empty()); + } + ASSERT_EQ( + version.writing_segment_meta()->writing_forward_block()->min_doc_id(), + 1u); + + // Pin the new writer's format before constructing a historical WAL. + std::vector docs; + ASSERT_NO_FATAL_FAILURE(ReadWalDocuments(path_, &docs)); + ASSERT_EQ(docs.size(), 2u); + EXPECT_EQ(docs[0]->get_operator(), Operator::UPDATE); + EXPECT_EQ(docs[0]->doc_id(), 0u); + EXPECT_EQ(docs[1]->get_operator(), Operator::UPDATE); + EXPECT_EQ(docs[1]->doc_id(), 1u); + ASSERT_NO_FATAL_FAILURE(RewriteWalAsLegacyUpserts(path_)); + + auto map = + IDMap::CreateAndOpen("wal_recovery", FindIdMap(path_), false, false); + ASSERT_NE(map, nullptr); + // The map can reach disk ahead of the manifest's deletion snapshot. + // Its latest replay ID no longer identifies committed predecessor ID 0. + ASSERT_TRUE(map->upsert("target", 2).ok()); + ASSERT_TRUE(map->flush().ok()); + } + + // Neither child flushes its recovered deletion bitmap. Both retries must + // rediscover ID 0, including when it belongs to another persisted segment. + for (int attempt = 0; attempt < 2; ++attempt) { + ASSERT_EXIT( + { + auto opened = Collection::Open(path_, CollectionOptions{}); + if (!opened.has_value()) { + std::cerr << opened.error() << std::endl; + std::_Exit(1); + } + ExpectOnlyTarget(opened.value(), "second"); + std::_Exit(::testing::Test::HasFailure() ? 2 : 0); + }, + ::testing::ExitedWithCode(0), ""); + } + + { + auto opened = Collection::Open(path_, CollectionOptions{}); + ASSERT_TRUE(opened.has_value()) << opened.error().message(); + ASSERT_NO_FATAL_FAILURE(ExpectOnlyTarget(opened.value(), "second")); + ASSERT_TRUE(opened.value()->flush().ok()); + } + CollectionOptions options; + options.read_only_ = true; + auto reopened = Collection::Open(path_, options); + ASSERT_TRUE(reopened.has_value()) << reopened.error().message(); + ASSERT_NO_FATAL_FAILURE(ExpectOnlyTarget(reopened.value(), "second")); + auto iterator = reopened.value()->create_iterator(); + ASSERT_TRUE(iterator.has_value()) << iterator.error().message(); + auto first = iterator.value()->next(); + ASSERT_TRUE(first.has_value()); + ASSERT_NE(first.value(), nullptr); + EXPECT_EQ(first.value()->pk_ref(), "target"); + EXPECT_EQ(first.value()->get("text"), "second"); + auto end = iterator.value()->next(); + ASSERT_TRUE(end.has_value()); + EXPECT_EQ(end.value(), nullptr); + } + + const std::string path_{"wal_recovery_regression_db"}; +}; + +TEST_F(WalRecoveryDeathTest, LargeStringAndTrailingDocRecoverFromWal) { + const std::vector ids{"prefix", "large", "suffix"}; + // Exceed the former reader-only limit using an ordinary short ID. + const std::vector values{ + "before", std::string(4 * 1024 * 1024 + 1024, 'v'), "after"}; + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(WriteStringDocsAndExit(path_, ids, values), + ::testing::ExitedWithCode(0), ""); + for (int reopen = 0; reopen < 2; ++reopen) { + auto opened = Collection::Open(path_, CollectionOptions{}); + ASSERT_TRUE(opened.has_value()) << opened.error().message(); + auto collection = std::move(opened).value(); + auto fetched = collection->fetch(ids); + ASSERT_TRUE(fetched.has_value()) << fetched.error().message(); + for (size_t i = 0; i < ids.size(); ++i) { + auto found = fetched.value().find(ids[i]); + ASSERT_NE(found, fetched.value().end()); + ASSERT_NE(found->second, nullptr); + EXPECT_EQ(found->second->pk_ref(), ids[i]); + EXPECT_EQ(found->second->get("text"), values[i]); + } + EXPECT_EQ(collection->stats().value().doc_count, ids.size()); + ASSERT_TRUE(collection->flush().ok()); + } +} + +TEST_F(WalRecoveryDeathTest, IncompleteTailAllowsLaterCrashRecovery) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT( + WriteStringDocsAndExit(path_, {"prefix", "torn"}, {"before", "tail"}), + ::testing::ExitedWithCode(0), ""); + if (!::testing::internal::InDeathTestChild()) { + const auto wal_path = FindWal(path_); + ASSERT_FALSE(wal_path.empty()); + std::filesystem::resize_file(wal_path, + std::filesystem::file_size(wal_path) - 1); + const auto tail_bytes = ReadFileBytes(wal_path); + const auto manifests = ReadManifests(path_); + ASSERT_FALSE(manifests.empty()); + { + CollectionOptions options; + options.read_only_ = true; + auto opened = Collection::Open(path_, options); + // Recovery needs writable index stores. A read-only attempt must fail + // explicitly and leave the WAL and committed manifest untouched. + ASSERT_FALSE(opened.has_value()); + EXPECT_EQ(opened.error().code(), StatusCode::FAILED_PRECONDITION); + EXPECT_NE(opened.error().message().find("read-write mode once"), + std::string::npos); + } + EXPECT_EQ(ReadFileBytes(wal_path), tail_bytes); + EXPECT_EQ(ReadManifests(path_), manifests); + } + + ASSERT_EXIT( + { + auto opened = Collection::Open(path_, CollectionOptions{}); + if (!opened.has_value()) { + std::cerr << opened.error() << std::endl; + std::_Exit(1); + } + Doc doc; + doc.set_pk("suffix"); + doc.set("text", "after"); + std::vector docs{doc}; + auto inserted = opened.value()->insert(docs); + if (!inserted.has_value() || !inserted.value().front().ok()) + std::_Exit(2); + std::_Exit(0); + }, + ::testing::ExitedWithCode(0), ""); + auto opened = Collection::Open(path_, CollectionOptions{}); + ASSERT_TRUE(opened.has_value()) << opened.error().message(); + auto fetched = opened.value()->fetch({"prefix", "torn", "suffix"}); + ASSERT_TRUE(fetched.has_value()); + ASSERT_NE(fetched.value().at("prefix"), nullptr); + ASSERT_NE(fetched.value().at("suffix"), nullptr); + EXPECT_TRUE(fetched.value().find("torn") == fetched.value().end() || + fetched.value().at("torn") == nullptr); + EXPECT_EQ(opened.value()->stats().value().doc_count, 2); +} + +TEST_F(WalRecoveryDeathTest, CorruptWalFailsOpenWithoutReplacingFiles) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT( + WriteStringDocsAndExit(path_, {"prefix", "broken"}, {"before", "after"}), + ::testing::ExitedWithCode(0), ""); + const auto wal_path = FindWal(path_); + ASSERT_FALSE(wal_path.empty()); + { + std::fstream file(wal_path, + std::ios::in | std::ios::out | std::ios::binary); + ASSERT_TRUE(file.is_open()); + uint32_t first_length; + file.seekg(64); + file.read(reinterpret_cast(&first_length), sizeof(first_length)); + ASSERT_TRUE(file.good()); + // Damage the second record CRC, keeping its complete framing intact. + file.seekp(64 + 8 + first_length + 4); + const uint32_t bad_crc = 0; + file.write(reinterpret_cast(&bad_crc), sizeof(bad_crc)); + ASSERT_TRUE(file.good()); + } + const auto bytes = ReadFileBytes(wal_path); + const auto manifests = ReadManifests(path_); + ASSERT_FALSE(manifests.empty()); + for (int attempt = 0; attempt < 2; ++attempt) { + auto opened = Collection::Open(path_, CollectionOptions{}); + ASSERT_FALSE(opened.has_value()); + EXPECT_NE(opened.error().message().find("CRC mismatch"), std::string::npos); + EXPECT_EQ(ReadFileBytes(wal_path), bytes); + EXPECT_EQ(ReadManifests(path_), manifests); + } +} + +TEST_F(WalRecoveryDeathTest, CorruptTailDoesNotApplyUpsertPrefix) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(WriteUpsertsAndExit(path_, true, true), + ::testing::ExitedWithCode(0), ""); + ASSERT_NO_FATAL_FAILURE(RewriteWalAsLegacyUpserts(path_)); + const auto wal_path = FindWal(path_); + ASSERT_FALSE(wal_path.empty()); + auto bytes = ReadFileBytes(wal_path); + size_t prefix_end = 64; + for (int i = 0; i < 2; ++i) { + uint32_t length; + ASSERT_GE(bytes.size() - prefix_end, 8u); + std::memcpy(&length, bytes.data() + prefix_end, sizeof(length)); + prefix_end += 8 + length; + ASSERT_LE(prefix_end, bytes.size()); + } + ASSERT_GE(bytes.size() - prefix_end, 8u); + bytes[prefix_end + 4] ^= 1; // Corrupt only the trailing record's CRC. + { + std::ofstream file(wal_path, std::ios::binary | std::ios::trunc); + file.write(bytes.data(), bytes.size()); + ASSERT_TRUE(file.good()); + } + const auto idmap_path = FindIdMap(path_); + ASSERT_FALSE(idmap_path.empty()); + { + auto map = IDMap::CreateAndOpen("wal_recovery", idmap_path, false, false); + ASSERT_NE(map, nullptr); + // The original row is committed as ID 0. Make the pre-recovery mapping + // deterministic even if RocksDB flushed uncommitted writes before exit. + ASSERT_TRUE(map->upsert("target", 0).ok()); + map->remove("broken"); + ASSERT_TRUE(map->flush().ok()); + } + const auto manifests = ReadManifests(path_); + for (int attempt = 0; attempt < 2; ++attempt) { + auto opened = Collection::Open(path_, CollectionOptions{}); + ASSERT_FALSE(opened.has_value()); + EXPECT_NE(opened.error().message().find("CRC mismatch"), std::string::npos); + EXPECT_EQ(ReadFileBytes(wal_path), bytes); + EXPECT_EQ(ReadManifests(path_), manifests); + auto map = IDMap::CreateAndOpen("wal_recovery", idmap_path, false, true); + ASSERT_NE(map, nullptr); + uint64_t original_id; + ASSERT_TRUE(map->has("target", &original_id)); + EXPECT_EQ(original_id, 0u); + EXPECT_FALSE(map->has("broken")); + } + + // Remove the damaged last record and retry the intact UPSERT prefix. + std::filesystem::resize_file(wal_path, prefix_end); + { + auto opened = Collection::Open(path_, CollectionOptions{}); + ASSERT_TRUE(opened.has_value()) << opened.error().message(); + auto fetched = opened.value()->fetch({"target"}); + ASSERT_TRUE(fetched.has_value()); + ASSERT_NE(fetched.value().at("target"), nullptr); + EXPECT_EQ(fetched.value().at("target")->get("text"), "second"); + EXPECT_EQ(opened.value()->stats().value().doc_count, 1); + Doc update; + update.set_pk("target"); + update.set("text", "third"); + std::vector updates{update}; + auto updated = opened.value()->upsert(updates); + ASSERT_TRUE(updated.has_value()); + ASSERT_TRUE(updated.value().front().ok()); + ASSERT_TRUE(opened.value()->flush().ok()); + } + auto reopened = Collection::Open(path_, CollectionOptions{}); + ASSERT_TRUE(reopened.has_value()); + auto fetched = reopened.value()->fetch({"target"}); + ASSERT_TRUE(fetched.has_value()); + ASSERT_NE(fetched.value().at("target"), nullptr); + EXPECT_EQ(fetched.value().at("target")->get("text"), "third"); + EXPECT_EQ(reopened.value()->stats().value().doc_count, 1); +} + +TEST_F(WalRecoveryDeathTest, UpsertReplayIgnoresSameAndLaterReplayIds) { + ::testing::FLAGS_gtest_death_test_style = "threadsafe"; + ASSERT_EXIT(WriteUpsertsAndExit(path_, false, false), + ::testing::ExitedWithCode(0), ""); + if (!::testing::internal::InDeathTestChild()) { + ASSERT_NO_FATAL_FAILURE(RewriteWalAsLegacyUpserts(path_)); + const auto idmap_path = FindIdMap(path_); + ASSERT_FALSE(idmap_path.empty()); + auto map = IDMap::CreateAndOpen("wal_recovery", idmap_path, false, false); + ASSERT_NE(map, nullptr); + // Simulate a partial previous recovery that persisted the second UPSERT's + // ID. It is later than the first replay ID and equal to the second. + ASSERT_TRUE(map->upsert("target", 1).ok()); + ASSERT_TRUE(map->flush().ok()); + } + ASSERT_EXIT( + { + auto opened = Collection::Open(path_, CollectionOptions{}); + if (!opened.has_value()) { + std::cerr << opened.error() << std::endl; + std::_Exit(1); + } + auto fetched = opened.value()->fetch({"target"}); + if (!fetched.has_value() || !fetched.value().at("target") || + fetched.value().at("target")->get("text") != + "second" || + opened.value()->stats().value().doc_count != 1) + std::_Exit(2); + Doc update; + update.set_pk("target"); + update.set("text", "third"); + std::vector updates{update}; + auto updated = opened.value()->upsert(updates); + if (!updated.has_value() || !updated.value().front().ok()) + std::_Exit(3); + std::_Exit(0); + }, + ::testing::ExitedWithCode(0), ""); + auto reopened = Collection::Open(path_, CollectionOptions{}); + ASSERT_TRUE(reopened.has_value()) << reopened.error().message(); + auto fetched = reopened.value()->fetch({"target"}); + ASSERT_TRUE(fetched.has_value()); + ASSERT_NE(fetched.value().at("target"), nullptr); + EXPECT_EQ(fetched.value().at("target")->get("text"), "third"); + EXPECT_EQ(reopened.value()->stats().value().doc_count, 1); +} + +TEST_F(WalRecoveryDeathTest, + LegacyUpsertsRecoverCommittedPredecessorInWritingSegment) { + ASSERT_NO_FATAL_FAILURE(CheckLegacyCommittedPredecessor(false)); +} + +TEST_F(WalRecoveryDeathTest, + LegacyUpsertsRecoverCommittedPredecessorInPersistedSegment) { + ASSERT_NO_FATAL_FAILURE(CheckLegacyCommittedPredecessor(true)); +} + +TEST_F(WalRecoveryDeathTest, UpsertWalRecordsInsertOrUpdatePredecessor) { + ASSERT_EXIT(WriteUpsertsAndExit(path_, false, false), + ::testing::ExitedWithCode(0), ""); + std::vector docs; + ASSERT_NO_FATAL_FAILURE(ReadWalDocuments(path_, &docs)); + ASSERT_EQ(docs.size(), 2u); + EXPECT_EQ(docs[0]->get_operator(), Operator::INSERT); + EXPECT_EQ(docs[0]->pk_ref(), "target"); + EXPECT_EQ(docs[0]->get("text"), "first"); + EXPECT_EQ(docs[1]->get_operator(), Operator::UPDATE); + EXPECT_EQ(docs[1]->doc_id(), 0u); + EXPECT_EQ(docs[1]->get("text"), "second"); +} + +} // namespace +} // namespace zvec diff --git a/tests/db/index/storage/wal_file_test.cc b/tests/db/index/storage/wal_file_test.cc index 50cd122da..a29e80e3d 100644 --- a/tests/db/index/storage/wal_file_test.cc +++ b/tests/db/index/storage/wal_file_test.cc @@ -12,20 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. -#ifdef _MSC_VER -#define _ALLOW_KEYWORD_MACROS -#endif -#define private public -#define protected public #include "db/index/storage/wal/wal_file.h" -#undef private -#undef protected - #include #include #include #include +#include #include +#include #include #include #include @@ -47,6 +41,16 @@ class WalFileTest : public testing::Test { } void TearDown() override {} + + // Legacy success-path tests use empty string as their loop sentinel, but + // still assert that the new API did not report an error. + std::string ReadRecord(const WalFilePtr &wal_file) { + auto result = wal_file->next(); + EXPECT_TRUE(result.has_value()) + << (result.has_value() ? "" : result.error().message()); + if (!result.has_value() || !result.value().has_value()) return {}; + return std::move(result.value().value()); + } }; TEST_F(WalFileTest, TestGeneral) { @@ -126,14 +130,14 @@ TEST_F(WalFileTest, TestGeneral) { uint32_t idx = 0; ret = wal_file->prepare_for_read(); ASSERT_EQ(ret, 0); - std::string record = wal_file->next(); + std::string record = ReadRecord(wal_file); while (!record.empty()) { if (idx < 100) { ASSERT_EQ(record, "hello"); } else { ASSERT_EQ(record, std::string("hello") + std::to_string(idx)); } - record = wal_file->next(); + record = ReadRecord(wal_file); idx++; } ASSERT_EQ(idx, 400); @@ -205,9 +209,9 @@ TEST_F(WalFileTest, TestMultiThread) { uint32_t idx = 0; ret = wal_file->prepare_for_read(); ASSERT_EQ(ret, 0); - std::string record = wal_file->next(); + std::string record = ReadRecord(wal_file); while (!record.empty()) { - record = wal_file->next(); + record = ReadRecord(wal_file); idx++; } ASSERT_EQ(idx, 30000); @@ -243,9 +247,9 @@ TEST_F(WalFileTest, TestBoundaryCondition) { ret = wal_file->open(wal_option); ASSERT_EQ(ret, 0); uint32_t idx = 0; - std::string record = wal_file->next(); + std::string record = ReadRecord(wal_file); while (!record.empty()) { - record = wal_file->next(); + record = ReadRecord(wal_file); idx++; } ASSERT_EQ(idx, 0); @@ -271,13 +275,13 @@ TEST_F(WalFileTest, TestBoundaryCondition) { idx = 0; ret = wal_file->prepare_for_read(); ASSERT_EQ(ret, 0); - record = wal_file->next(); + record = ReadRecord(wal_file); while (!record.empty()) { ASSERT_EQ(record.size(), 4); for (size_t i = 0; i < 4; i++) { ASSERT_EQ(record[i], i); } - record = wal_file->next(); + record = ReadRecord(wal_file); idx++; } ASSERT_EQ(idx, 1); @@ -312,13 +316,13 @@ TEST_F(WalFileTest, TestBoundaryCondition) { idx = 0; ret = wal_file->prepare_for_read(); ASSERT_EQ(ret, 0); - record = wal_file->next(); + record = ReadRecord(wal_file); while (!record.empty()) { ASSERT_EQ(record.size(), BIG_DATA_SIZE); for (size_t i = 0; i < BIG_DATA_SIZE; i++) { ASSERT_EQ((uint8_t)record[i], i % 256); } - record = wal_file->next(); + record = ReadRecord(wal_file); idx++; } ASSERT_EQ(idx, 1); @@ -349,10 +353,10 @@ TEST_F(WalFileTest, TestBoundaryCondition) { idx = 0; ret = wal_file->prepare_for_read(); ASSERT_EQ(ret, 0); - record = wal_file->next(); + record = ReadRecord(wal_file); while (!record.empty()) { ASSERT_EQ(record, std::string("hello") + std::to_string(idx)); - record = wal_file->next(); + record = ReadRecord(wal_file); idx++; } ASSERT_EQ(idx, 99); @@ -417,16 +421,15 @@ TEST_F(WalFileTest, TestFirstErrorCase) { ret = wal_file->open(wal_option); ASSERT_EQ(ret, 0); - uint32_t idx = 0; - ret = wal_file->prepare_for_read(); - ASSERT_EQ(ret, 0); - std::string record = wal_file->next(); - while (!record.empty()) { - ASSERT_EQ(record, "hello"); - record = wal_file->next(); - idx++; + ASSERT_EQ(wal_file->prepare_for_read(), 0); + for (size_t i = 0; i < 0; ++i) { + EXPECT_EQ(ReadRecord(wal_file), "hello"); } - ASSERT_EQ(idx, 0); + auto result = wal_file->next(); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), StatusCode::INTERNAL_ERROR); + EXPECT_NE(result.error().message().find("CRC mismatch"), std::string::npos); + EXPECT_NE(wal_file->append("after corruption"), 0); // close ret = wal_file->close(); ASSERT_EQ(ret, 0); @@ -477,16 +480,15 @@ TEST_F(WalFileTest, TestMiddleErrorCase) { ret = wal_file->open(wal_option); ASSERT_EQ(ret, 0); - uint32_t idx = 0; - ret = wal_file->prepare_for_read(); - ASSERT_EQ(ret, 0); - std::string record = wal_file->next(); - while (!record.empty()) { - ASSERT_EQ(record, "hello"); - record = wal_file->next(); - idx++; + ASSERT_EQ(wal_file->prepare_for_read(), 0); + for (size_t i = 0; i < 5; ++i) { + EXPECT_EQ(ReadRecord(wal_file), "hello"); } - ASSERT_EQ(idx, 5); + auto result = wal_file->next(); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), StatusCode::INTERNAL_ERROR); + EXPECT_NE(result.error().message().find("CRC mismatch"), std::string::npos); + EXPECT_NE(wal_file->append("after corruption"), 0); // close ret = wal_file->close(); ASSERT_EQ(ret, 0); @@ -538,10 +540,10 @@ TEST_F(WalFileTest, TestLastErrorCase) { uint32_t idx = 0; ret = wal_file->prepare_for_read(); ASSERT_EQ(ret, 0); - std::string record = wal_file->next(); + std::string record = ReadRecord(wal_file); while (!record.empty()) { ASSERT_EQ(record, "hello"); - record = wal_file->next(); + record = ReadRecord(wal_file); idx++; } ASSERT_EQ(idx, 9); @@ -594,16 +596,15 @@ TEST_F(WalFileTest, TestLengthSmallErrorCase) { ret = wal_file->open(wal_option); ASSERT_EQ(ret, 0); - uint32_t idx = 0; - ret = wal_file->prepare_for_read(); - ASSERT_EQ(ret, 0); - std::string record = wal_file->next(); - while (!record.empty()) { - ASSERT_EQ(record, "hello"); - record = wal_file->next(); - idx++; + ASSERT_EQ(wal_file->prepare_for_read(), 0); + for (size_t i = 0; i < 0; ++i) { + EXPECT_EQ(ReadRecord(wal_file), "hello"); } - ASSERT_EQ(idx, 0); + auto result = wal_file->next(); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), StatusCode::INTERNAL_ERROR); + EXPECT_NE(result.error().message().find("CRC mismatch"), std::string::npos); + EXPECT_NE(wal_file->append("after corruption"), 0); // close ret = wal_file->close(); ASSERT_EQ(ret, 0); @@ -643,7 +644,7 @@ TEST_F(WalFileTest, TestLengthBigErrorCase) { dir_path, "data.wal.", std::to_string(segment_id)); int wal_fd = open(wal_path.c_str(), O_RDWR, 0644); ASSERT_GT(wal_fd, 0); - uint32_t err_length = 200; // exceed file size 130 + uint32_t err_length = std::numeric_limits::max(); lseek(wal_fd, 64, SEEK_SET); write(wal_fd, (const void *)&err_length, 4); @@ -657,10 +658,10 @@ TEST_F(WalFileTest, TestLengthBigErrorCase) { uint32_t idx = 0; ret = wal_file->prepare_for_read(); ASSERT_EQ(ret, 0); - std::string record = wal_file->next(); + std::string record = ReadRecord(wal_file); while (!record.empty()) { ASSERT_EQ(record, "hello"); - record = wal_file->next(); + record = ReadRecord(wal_file); idx++; } ASSERT_EQ(idx, 0); @@ -714,16 +715,15 @@ TEST_F(WalFileTest, TestCRCErrorCase) { ret = wal_file->open(wal_option); ASSERT_EQ(ret, 0); - uint32_t idx = 0; - ret = wal_file->prepare_for_read(); - ASSERT_EQ(ret, 0); - std::string record = wal_file->next(); - while (!record.empty()) { - ASSERT_EQ(record, "hello"); - record = wal_file->next(); - idx++; + ASSERT_EQ(wal_file->prepare_for_read(), 0); + for (size_t i = 0; i < 1; ++i) { + EXPECT_EQ(ReadRecord(wal_file), "hello"); } - ASSERT_EQ(idx, 1); + auto result = wal_file->next(); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), StatusCode::INTERNAL_ERROR); + EXPECT_NE(result.error().message().find("CRC mismatch"), std::string::npos); + EXPECT_NE(wal_file->append("after corruption"), 0); // close ret = wal_file->close(); ASSERT_EQ(ret, 0); @@ -732,6 +732,115 @@ TEST_F(WalFileTest, TestCRCErrorCase) { ASSERT_EQ(ret, 0); } +TEST_F(WalFileTest, RecordLargerThanFourMiBPreservesFollowingRecord) { + const std::string path = "./data.wal.large"; + auto wal = WalFile::Create(path); + WalOptions options; + options.create_new = true; + ASSERT_EQ(wal->open(options), 0); + const std::string large(4 * 1024 * 1024 + 1024, 'x'); + ASSERT_EQ(wal->append("prefix"), 0); + ASSERT_EQ(wal->append(std::string(large)), 0); + ASSERT_EQ(wal->append("suffix"), 0); + ASSERT_EQ(wal->close(), 0); + options.create_new = false; + ASSERT_EQ(wal->open(options), 0); + ASSERT_EQ(wal->prepare_for_read(), 0); + EXPECT_EQ(ReadRecord(wal), "prefix"); + EXPECT_EQ(ReadRecord(wal), large); + EXPECT_EQ(ReadRecord(wal), "suffix"); + auto end = wal->next(); + ASSERT_TRUE(end.has_value()); + EXPECT_FALSE(end.value().has_value()); +} + +TEST_F(WalFileTest, IncompleteTailIsRemovedOnlyBeforeAppend) { + const std::string path = "./data.wal.tail"; + constexpr size_t prefix_end = 64 + 8 + 6; + // Exercise every partial header and partial payload boundary. + for (size_t tail_size = 1; tail_size < 8 + 4; ++tail_size) { + SCOPED_TRACE(tail_size); + auto wal = WalFile::Create(path); + WalOptions options; + options.create_new = true; + ASSERT_EQ(wal->open(options), 0); + ASSERT_EQ(wal->append("prefix"), 0); + ASSERT_EQ(wal->append("torn"), 0); + ASSERT_EQ(wal->close(), 0); + ailego::File file; + ASSERT_TRUE(file.open(path, false)); + ASSERT_TRUE(file.truncate(prefix_end + tail_size)); + file.close(); + + options.create_new = false; + ASSERT_EQ(wal->open(options), 0); + ASSERT_EQ(wal->prepare_for_read(), 0); + EXPECT_EQ(ReadRecord(wal), "prefix"); + auto end = wal->next(); + ASSERT_TRUE(end.has_value()); + EXPECT_FALSE(end.value().has_value()); + ASSERT_TRUE(file.open(path, true)); + EXPECT_EQ(file.size(), prefix_end + tail_size); + file.close(); + + ASSERT_EQ(wal->append("suffix"), 0); + ASSERT_EQ(wal->close(), 0); + ASSERT_EQ(wal->open(options), 0); + ASSERT_EQ(wal->prepare_for_read(), 0); + EXPECT_EQ(ReadRecord(wal), "prefix"); + EXPECT_EQ(ReadRecord(wal), "suffix"); + end = wal->next(); + ASSERT_TRUE(end.has_value()); + EXPECT_FALSE(end.value().has_value()); + ASSERT_EQ(wal->remove(), 0); + } +} + +TEST_F(WalFileTest, ClosedReaderReturnsError) { + auto wal = WalFile::Create("./data.wal.closed"); + auto result = wal->next(); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().code(), StatusCode::INTERNAL_ERROR); +} + +TEST_F(WalFileTest, ZeroLengthRecordIsCorruption) { + const std::string path = "./data.wal.zero"; + auto wal = WalFile::Create(path); + WalOptions options; + options.create_new = true; + ASSERT_EQ(wal->open(options), 0); + EXPECT_NE(wal->append(""), 0); + ASSERT_EQ(wal->append("payload"), 0); + ASSERT_EQ(wal->close(), 0); + ailego::File file; + ASSERT_TRUE(file.open(path, false)); + const uint32_t length = 0; + ASSERT_EQ(file.write(64, &length, sizeof(length)), sizeof(length)); + file.close(); + options.create_new = false; + ASSERT_EQ(wal->open(options), 0); + ASSERT_EQ(wal->prepare_for_read(), 0); + auto result = wal->next(); + ASSERT_FALSE(result.has_value()); + EXPECT_NE(result.error().message().find("zero length"), std::string::npos); +} + +TEST_F(WalFileTest, TruncatedFileHeaderReturnsError) { + const std::string path = "./data.wal.header"; + auto wal = WalFile::Create(path); + WalOptions options; + options.create_new = true; + ASSERT_EQ(wal->open(options), 0); + ASSERT_EQ(wal->close(), 0); + ailego::File file; + ASSERT_TRUE(file.open(path, false)); + ASSERT_TRUE(file.truncate(63)); + file.close(); + options.create_new = false; + ASSERT_EQ(wal->open(options), 0); + EXPECT_NE(wal->prepare_for_read(), 0); +} + #if defined(__GNUC__) || defined(__GNUG__) #pragma GCC diagnostic pop #endif \ No newline at end of file From ee3e3140b6e5a5e420bd39d89ad690fbad187ef1 Mon Sep 17 00:00:00 2001 From: Qinren Zhou Date: Fri, 18 Sep 2026 00:37:55 +0800 Subject: [PATCH 2/2] fix: preserve WAL writes across block rotation and validate recovery --- src/db/index/common/doc.cc | 505 ++++++++----------- src/db/index/segment/segment.cc | 48 +- src/db/index/storage/wal/local_wal_file.cc | 36 +- src/db/index/storage/wal/local_wal_file.h | 8 +- src/db/index/storage/wal/wal_file.h | 2 + src/include/zvec/db/doc.h | 5 - tests/db/crash_recovery/wal_recovery_test.cc | 127 +++++ tests/db/index/common/doc_test.cc | 92 ++++ tests/db/index/storage/wal_file_test.cc | 57 ++- 9 files changed, 562 insertions(+), 318 deletions(-) diff --git a/src/db/index/common/doc.cc b/src/db/index/common/doc.cc index 48742ba2f..607cd2849 100644 --- a/src/db/index/common/doc.cc +++ b/src/db/index/common/doc.cc @@ -122,28 +122,11 @@ namespace { template T byte_swap(T value) { - if constexpr (std::is_same_v) { - uint16_t val; - std::memcpy(&val, static_cast(&value), sizeof(val)); - val = ailego_bswap16(val); - float16_t result; - std::memcpy(static_cast(&result), &val, sizeof(result)); - return result; - } else if constexpr (sizeof(T) == 1) { - return value; - } else if constexpr (sizeof(T) == 2) { - return (value << 8) | ((value >> 8) & 0xFF); - } else if constexpr (sizeof(T) == 4) { - return static_cast(ailego_bswap32(static_cast(value))); - } else if constexpr (sizeof(T) == 8) { - return static_cast(ailego_bswap64(static_cast(value))); - } else { - T result = 0; - for (size_t i = 0; i < sizeof(T); ++i) { - result |= ((value >> (i * 8)) & 0xFF) << ((sizeof(T) - 1 - i) * 8); - } - return result; - } + T result; + const auto *source = reinterpret_cast(&value); + auto *destination = reinterpret_cast(&result); + std::reverse_copy(source, source + sizeof(T), destination); + return result; } template @@ -156,17 +139,179 @@ void write_value_to_buffer(std::vector &buffer, const T &value) { buffer.insert(buffer.end(), bytes, bytes + sizeof(T)); } -template -T read_value_from_buffer(const uint8_t *&data) { - T value; - std::memcpy(&value, data, sizeof(T)); - data += sizeof(T); +// Read persisted values only after checking their complete byte range. Length +// fields are checked before allocation, including each nested array/string. +class DocBufferReader { + public: + DocBufferReader(const uint8_t *data, size_t size) + : data_(data), remaining_(size) {} - if (IS_BIG_ENDIAN) { - value = byte_swap(value); + size_t remaining() const { + return remaining_; + } + + bool read_bytes(void *destination, size_t size) { + if (size > remaining_) return false; + if (size != 0) { + std::memcpy(destination, data_, size); + data_ += size; + remaining_ -= size; + } + return true; + } + + template + bool read_native(T &value) { + return read_bytes(&value, sizeof(T)); + } + + bool read_string_bytes(std::string &value, size_t size) { + if (size > remaining_) return false; + value.assign(reinterpret_cast(data_), size); + data_ += size; + remaining_ -= size; + return true; + } + + bool read_value(Doc::Value &value) { + uint8_t type; + if (!read_native(type)) return false; + switch (type) { + case TYPE_EMPTY: + value = std::monostate{}; + return true; + case TYPE_BOOL: + return read_as(value); + case TYPE_INT32: + return read_as(value); + case TYPE_UINT32: + return read_as(value); + case TYPE_INT64: + return read_as(value); + case TYPE_UINT64: + return read_as(value); + case TYPE_FLOAT: + return read_as(value); + case TYPE_DOUBLE: + return read_as(value); + case TYPE_STRING: + return read_as(value); + case TYPE_VECTOR_BOOL: + return read_as>(value); + case TYPE_VECTOR_INT8: + return read_as>(value); + case TYPE_VECTOR_INT16: + return read_as>(value); + case TYPE_VECTOR_INT32: + return read_as>(value); + case TYPE_VECTOR_INT64: + return read_as>(value); + case TYPE_VECTOR_UINT32: + return read_as>(value); + case TYPE_VECTOR_UINT64: + return read_as>(value); + case TYPE_VECTOR_FLOAT16: + return read_as>(value); + case TYPE_VECTOR_FLOAT: + return read_as>(value); + case TYPE_VECTOR_DOUBLE: + return read_as>(value); + case TYPE_VECTOR_STRING: + return read_as>(value); + case TYPE_VECTOR_PAIR_INT_FLOAT: + return read_as, std::vector>>( + value); + case TYPE_VECTOR_PAIR_INT_FLOAT16: + return read_as< + std::pair, std::vector>>(value); + default: + return false; + } + } + + private: + template + bool read_little_endian(T &value) { + if (!read_native(value)) return false; + if (IS_BIG_ENDIAN) { + auto *bytes = reinterpret_cast(&value); + std::reverse(bytes, bytes + sizeof(T)); + } + return true; + } + + template + bool read_as(Doc::Value &out) { + T value; + if (!read(value)) return false; + out = std::move(value); + return true; + } + + template + bool read(T &value) { + return read_little_endian(value); + } + + bool read(bool &value) { + static_assert(sizeof(bool) == sizeof(uint8_t)); + uint8_t byte; + if (!read_native(byte) || byte > 1) return false; + value = byte != 0; + return true; } - return value; -} + + bool read(std::string &value) { + uint32_t size; + return read_little_endian(size) && read_string_bytes(value, size); + } + + template + bool read(std::vector &values) { + uint32_t count; + if (!read_little_endian(count)) return false; + if constexpr (std::is_same_v) { + // Each string contains at least its four-byte length prefix. + if (count > remaining_ / sizeof(uint32_t)) return false; + values.reserve(count); + for (uint32_t i = 0; i < count; ++i) { + std::string value; + if (!read(value)) return false; + values.push_back(std::move(value)); + } + } else if constexpr (std::is_same_v) { + if (count > remaining_ / sizeof(bool)) return false; + values.reserve(count); + for (uint32_t i = 0; i < count; ++i) { + bool value; + if (!read(value)) return false; + values.push_back(value); + } + } else { + // Division avoids overflow before checking the allocation/copy size. + if (count > remaining_ / sizeof(T)) return false; + values.resize(count); + if (!read_bytes(values.data(), static_cast(count) * sizeof(T))) { + return false; + } + if (IS_BIG_ENDIAN) { + for (auto &value : values) { + auto *bytes = reinterpret_cast(&value); + std::reverse(bytes, bytes + sizeof(T)); + } + } + } + return true; + } + + template + bool read(std::pair, std::vector> &value) { + return read(value.first) && read(value.second); + } + + const uint8_t *data_; + size_t remaining_; +}; template std::string vec_to_string(const std::vector &v) { @@ -198,11 +343,6 @@ void Doc::write_to_buffer(std::vector &buffer, const void *src, buffer.insert(buffer.end(), bytes, bytes + size); } -void Doc::read_from_buffer(const uint8_t *&data, void *dest, size_t size) { - std::memcpy(dest, data, size); - data += size; -} - void Doc::serialize_value(std::vector &buffer, const Value &value) { std::visit( [&buffer](const auto &v) { @@ -436,238 +576,6 @@ void Doc::serialize_value(std::vector &buffer, const Value &value) { } -Doc::Value Doc::deserialize_value(const uint8_t *&data) { - uint8_t type; - read_from_buffer(data, &type, sizeof(type)); - - switch (type) { - case TYPE_EMPTY: { - return std::monostate{}; - } - case TYPE_BOOL: { - bool v; - read_from_buffer(data, &v, sizeof(v)); - return v; - } - case TYPE_INT32: { - return read_value_from_buffer(data); - } - case TYPE_INT64: { - return read_value_from_buffer(data); - } - case TYPE_UINT32: { - return read_value_from_buffer(data); - } - case TYPE_UINT64: { - return read_value_from_buffer(data); - } - case TYPE_FLOAT: { - return read_value_from_buffer(data); - } - case TYPE_DOUBLE: { - return read_value_from_buffer(data); - } - case TYPE_STRING: { - uint32_t len = read_value_from_buffer(data); - std::string v(reinterpret_cast(data), len); - data += len; - return v; - } - case TYPE_VECTOR_BOOL: { - uint32_t len = read_value_from_buffer(data); - std::vector v; - v.reserve(len); - for (uint32_t i = 0; i < len; ++i) { - bool b; - read_from_buffer(data, &b, sizeof(b)); - v.push_back(b); - } - return v; - } - case TYPE_VECTOR_INT8: { - uint32_t len = read_value_from_buffer(data); - std::vector v(len); - read_from_buffer(data, v.data(), len * sizeof(int8_t)); - return v; - } - case TYPE_VECTOR_INT16: { - uint32_t len = read_value_from_buffer(data); - std::vector v(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v[i] = byte_swap(read_value_from_buffer(data)); - } - } else { - read_from_buffer(data, v.data(), len * sizeof(int16_t)); - } - return v; - } - case TYPE_VECTOR_INT32: { - uint32_t len = read_value_from_buffer(data); - std::vector v(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v[i] = byte_swap(read_value_from_buffer(data)); - } - } else { - read_from_buffer(data, v.data(), len * sizeof(int32_t)); - } - return v; - } - case TYPE_VECTOR_INT64: { - uint32_t len = read_value_from_buffer(data); - std::vector v(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v[i] = byte_swap(read_value_from_buffer(data)); - } - } else { - read_from_buffer(data, v.data(), len * sizeof(int64_t)); - } - return v; - } - case TYPE_VECTOR_UINT32: { - uint32_t len = read_value_from_buffer(data); - std::vector v(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v[i] = byte_swap(read_value_from_buffer(data)); - } - } else { - read_from_buffer(data, v.data(), len * sizeof(uint32_t)); - } - return v; - } - case TYPE_VECTOR_UINT64: { - uint32_t len = read_value_from_buffer(data); - std::vector v(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v[i] = byte_swap(read_value_from_buffer(data)); - } - } else { - read_from_buffer(data, v.data(), len * sizeof(uint64_t)); - } - return v; - } - case TYPE_VECTOR_FLOAT: { - uint32_t len = read_value_from_buffer(data); - std::vector v(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v[i] = byte_swap(read_value_from_buffer(data)); - } - } else { - read_from_buffer(data, v.data(), len * sizeof(float)); - } - return v; - } - case TYPE_VECTOR_DOUBLE: { - uint32_t len = read_value_from_buffer(data); - std::vector v(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v[i] = byte_swap(read_value_from_buffer(data)); - } - } else { - read_from_buffer(data, v.data(), len * sizeof(double)); - } - return v; - } - case TYPE_VECTOR_FLOAT16: { - uint32_t len = read_value_from_buffer(data); - std::vector v(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v[i] = byte_swap(read_value_from_buffer(data)); - } - } else { - read_from_buffer(data, v.data(), len * sizeof(float16_t)); - } - return v; - } - case TYPE_VECTOR_STRING: { - uint32_t len = read_value_from_buffer(data); - std::vector v; - v.reserve(len); - for (uint32_t i = 0; i < len; ++i) { - uint32_t str_len = read_value_from_buffer(data); - std::string s(reinterpret_cast(data), str_len); - data += str_len; - v.push_back(s); - } - return v; - } - case TYPE_VECTOR_PAIR_INT_FLOAT: { - uint32_t len = read_value_from_buffer(data); - std::pair, std::vector> v; - v.first.reserve(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v.first.push_back( - byte_swap(read_value_from_buffer(data))); - } - } else { - for (uint32_t i = 0; i < len; ++i) { - uint32_t first; - read_from_buffer(data, &first, sizeof(first)); - v.first.push_back(first); - } - } - len = read_value_from_buffer(data); - v.second.reserve(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v.second.push_back( - byte_swap(read_value_from_buffer(data))); - } - } else { - for (uint32_t i = 0; i < len; ++i) { - float second; - read_from_buffer(data, &second, sizeof(second)); - v.second.push_back(second); - } - } - return v; - } - case TYPE_VECTOR_PAIR_INT_FLOAT16: { - uint32_t len = read_value_from_buffer(data); - std::pair, std::vector> v; - v.first.reserve(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v.first.push_back( - byte_swap(read_value_from_buffer(data))); - } - } else { - for (uint32_t i = 0; i < len; ++i) { - uint32_t first; - read_from_buffer(data, &first, sizeof(first)); - v.first.push_back(first); - } - } - len = read_value_from_buffer(data); - v.second.reserve(len); - if (IS_BIG_ENDIAN) { - for (uint32_t i = 0; i < len; ++i) { - v.second.push_back( - byte_swap(read_value_from_buffer(data))); - } - } else { - for (uint32_t i = 0; i < len; ++i) { - float16_t second; - read_from_buffer(data, &second, sizeof(second)); - v.second.push_back(second); - } - } - return v; - } - - default: - throw std::runtime_error("Unknown value type: " + std::to_string(type)); - } -} - std::vector Doc::serialize() const { std::vector buffer; uint32_t pk_len = static_cast(pk_.size()); @@ -692,37 +600,40 @@ std::vector Doc::serialize() const { return buffer; } -Doc::Ptr Doc::deserialize(const uint8_t *data, size_t /*size*/) { - const uint8_t *ptr = data; - Doc::Ptr doc = std::make_shared(); - - uint32_t pk_len = read_value_from_buffer(ptr); - std::string pk(reinterpret_cast(ptr), pk_len); - ptr += pk_len; - doc->set_pk(pk); - - float score = read_value_from_buffer(ptr); - doc->set_score(score); - - uint64_t doc_id = read_value_from_buffer(ptr); - doc->set_doc_id(doc_id); - - Operator op; - read_from_buffer(ptr, &op, sizeof(op)); - doc->set_operator(op); - - uint32_t field_count = read_value_from_buffer(ptr); - +Doc::Ptr Doc::deserialize(const uint8_t *data, size_t size) { + if (!data) return nullptr; + DocBufferReader reader(data, size); + auto doc = std::make_shared(); + uint32_t pk_length; + uint32_t operation; + uint32_t field_count; + // The document header and field-name lengths retain their existing native + // representation; value payloads use the existing little-endian encoding. + if (!reader.read_native(pk_length) || + !reader.read_string_bytes(doc->pk_, pk_length) || + !reader.read_native(doc->score_) || !reader.read_native(doc->doc_id_) || + !reader.read_native(operation) || + operation > static_cast(Operator::DELETE) || + !reader.read_native(field_count)) { + return nullptr; + } + doc->op_ = static_cast(operation); + // Even an empty name and a null value require a length prefix and type byte. + if (field_count > reader.remaining() / (sizeof(uint32_t) + sizeof(uint8_t))) { + return nullptr; + } for (uint32_t i = 0; i < field_count; ++i) { - uint32_t name_len = read_value_from_buffer(ptr); - std::string field_name(reinterpret_cast(ptr), name_len); - ptr += name_len; - - Doc::Value value = deserialize_value(ptr); - doc->fields_[field_name] = value; + uint32_t name_length; + std::string name; + Value value; + if (!reader.read_native(name_length) || + !reader.read_string_bytes(name, name_length) || + !reader.read_value(value) || + !doc->fields_.emplace(std::move(name), std::move(value)).second) { + return nullptr; + } } - - return doc; + return reader.remaining() == 0 ? doc : nullptr; } Status Doc::validate_and_sanitize(const CollectionSchema::Ptr &schema, diff --git a/src/db/index/segment/segment.cc b/src/db/index/segment/segment.cc index 26742c8d7..a84bd142a 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -20,6 +20,8 @@ #include #include #include +#include +#include #include #include #include @@ -905,11 +907,6 @@ Status SegmentImpl::internal_insert(Doc &doc) { uint64_t g_doc_id = doc_id_allocator_.fetch_add(1); doc.set_doc_id(g_doc_id); - if (ready_for_dump_block()) { - auto s = flush(); - CHECK_RETURN_STATUS(s); - } - // init writing memory components if (!memory_store_) { auto s = init_memory_components(); @@ -4362,6 +4359,7 @@ Status SegmentImpl::recover() { WalFilePtr recover_wal_file; WalOptions wal_option; wal_option.create_new = false; + wal_option.read_only = options_.read_only_; if (WalFile::CreateAndOpen(wal_file_path, wal_option, &recover_wal_file) != 0) { LOG_ERROR("WAL recovery failed: unable to open WAL file [%s]", @@ -4388,6 +4386,7 @@ Status SegmentImpl::recover() { for (int pass = 0; pass < 2; ++pass) { const bool replay = pass == 1; total_recovered_doc_count = 0; + uint64_t next_doc_id = first_replay_id; int ret = recover_wal_file->prepare_for_read(); if (ret != 0) { LOG_ERROR( @@ -4428,8 +4427,17 @@ Status SegmentImpl::recover() { } const auto &buf = record.value().value(); total_recovered_doc_count++; - auto doc = Doc::deserialize(reinterpret_cast(buf.data()), - buf.size()); + Doc::Ptr doc; + try { + doc = Doc::deserialize(reinterpret_cast(buf.data()), + buf.size()); + } catch (const std::bad_alloc &) { + return Status(StatusCode::RESOURCE_EXHAUSTED, + "Unable to allocate WAL document during recovery"); + } catch (const std::length_error &) { + return Status::InternalError("Invalid WAL document length: ", + wal_file_path); + } if (doc == nullptr) { LOG_ERROR( "WAL record recovery failed: path[%s], segment[%d], record[%zu], " @@ -4440,6 +4448,23 @@ Status SegmentImpl::recover() { "], record[", total_recovered_doc_count, "]"); } + // Check references during the first pass as well: an UPDATE/DELETE may + // only target an earlier row, never a row yet to be allocated by replay. + const auto operation = doc->get_operator(); + if ((operation == Operator::UPDATE || operation == Operator::DELETE) && + doc->doc_id() >= next_doc_id) { + return Status::InternalError("Invalid WAL predecessor: path[", + wal_file_path, "], record[", + total_recovered_doc_count, "]"); + } + if (operation != Operator::DELETE) { + if (next_doc_id == std::numeric_limits::max()) { + return Status::InternalError("WAL document ID overflow: ", + wal_file_path); + } + ++next_doc_id; + } + if (!replay) { if (doc->get_operator() == Operator::UPSERT) { legacy_upsert_keys.insert(doc->pk_ref()); @@ -4545,6 +4570,15 @@ Status SegmentImpl::open_wal_file() { } Status SegmentImpl::append_wal(const Doc &doc) { + // Rotate a full block before logging the next operation. Flushing inside + // internal_insert would delete the WAL containing this still-unapplied + // operation (and could commit an UPDATE's deletion without its replacement). + // Recovery bypasses this path and keeps its source WAL until replay finishes. + if (ready_for_dump_block()) { + auto status = flush(); + CHECK_RETURN_STATUS(status); + } + std::vector buf = doc.serialize(); if (!wal_file_) { diff --git a/src/db/index/storage/wal/local_wal_file.cc b/src/db/index/storage/wal/local_wal_file.cc index e9d0493c5..a2d084a88 100644 --- a/src/db/index/storage/wal/local_wal_file.cc +++ b/src/db/index/storage/wal/local_wal_file.cc @@ -39,9 +39,12 @@ int LocalWalFile::append(std::string &&data) { record.content_ = std::move(data); std::lock_guard lock(file_mutex_); - if (!opened_ || failed_) { + if (!opened_ || failed_ || read_only_) { return -1; } + // Append changes the shared file position. Require a fresh read preparation + // instead of silently skipping unread records after an intervening write. + reader_ready_ = false; if (incomplete_tail_offset_) { if (!file_.truncate(*incomplete_tail_offset_)) { WLOG_ERROR("Wal incomplete tail truncation failed"); @@ -71,7 +74,7 @@ int LocalWalFile::append(std::string &&data) { Result> LocalWalFile::next() { std::lock_guard lock(file_mutex_); - if (!opened_ || failed_) { + if (!opened_ || failed_ || !reader_ready_) { return tl::make_unexpected( Status::InternalError("WAL is not open for reading or has failed")); } @@ -95,7 +98,9 @@ Result> LocalWalFile::next() { } int LocalWalFile::open(const WalOptions &wal_option) { + std::lock_guard lock(file_mutex_); CHECK_STATUS(opened_, false); + if (wal_option.create_new && wal_option.read_only) return -1; if (wal_option.create_new) { if (FileHelper::FileExists(wal_path_)) { WLOG_ERROR("Wal open error. file already exist create_new[%d]", @@ -109,10 +114,12 @@ int LocalWalFile::open(const WalOptions &wal_option) { } // write wal header + header_ = WalHeader{}; size_t write_size = file_.write((const void *)&header_, sizeof(header_)); if (write_size != sizeof(header_)) { WLOG_ERROR("Wal write header error. create_new[%d]", wal_option.create_new); + file_.close(); return -1; } @@ -123,19 +130,22 @@ int LocalWalFile::open(const WalOptions &wal_option) { return -1; } - if (!file_.open(wal_path_.c_str(), false)) { + if (!file_.open(wal_path_.c_str(), wal_option.read_only)) { WLOG_ERROR("Wal open error. create_new[%d]", wal_option.create_new); return -1; } // open default for write if (!file_.seek(0, ailego::File::Origin::End)) { + file_.close(); return -1; } } max_docs_wal_flush_ = wal_option.max_docs_wal_flush; opened_ = true; + read_only_ = wal_option.read_only; + reader_ready_ = false; failed_ = false; incomplete_tail_offset_.reset(); docs_count_ = 0; @@ -145,35 +155,50 @@ int LocalWalFile::open(const WalOptions &wal_option) { } int LocalWalFile::close() { + std::lock_guard lock(file_mutex_); CHECK_STATUS(opened_, true); file_.close(); WLOG_INFO("Wal close success"); opened_ = false; + reader_ready_ = false; return 0; } int LocalWalFile::remove() { + std::lock_guard lock(file_mutex_); + if (read_only_) return -1; if (opened_) { - close(); + file_.close(); + opened_ = false; + reader_ready_ = false; } if (FileHelper::FileExists(wal_path_)) { - FileHelper::RemoveFile(wal_path_); + if (!FileHelper::RemoveFile(wal_path_)) { + WLOG_ERROR("Wal remove failed"); + return -1; + } WLOG_INFO("Wal remove success."); } return 0; } int LocalWalFile::flush() { + std::lock_guard lock(file_mutex_); CHECK_STATUS(opened_, true); + if (failed_ || read_only_) return -1; if (!file_.flush()) { WLOG_ERROR("Wal flush error."); + failed_ = true; return -1; } + docs_count_ = 0; return 0; } int LocalWalFile::prepare_for_read() { + std::lock_guard lock(file_mutex_); CHECK_STATUS(opened_, true); + reader_ready_ = false; incomplete_tail_offset_.reset(); if (failed_ || !file_.seek(0, ailego::File::Origin::Begin)) { return -1; @@ -189,6 +214,7 @@ int LocalWalFile::prepare_for_read() { failed_ = true; return -1; } + reader_ready_ = true; return 0; } diff --git a/src/db/index/storage/wal/local_wal_file.h b/src/db/index/storage/wal/local_wal_file.h index c2a8dcadf..21c4dd987 100644 --- a/src/db/index/storage/wal/local_wal_file.h +++ b/src/db/index/storage/wal/local_wal_file.h @@ -13,7 +13,6 @@ // limitations under the License. #pragma once -#include #include #include #include "wal_file.h" @@ -68,7 +67,8 @@ class LocalWalFile : public WalFile { int remove() override; bool has_record() override { - return file_.size() > sizeof(header_); + std::lock_guard lock(file_mutex_); + return opened_ && file_.size() > sizeof(header_); } private: @@ -84,10 +84,12 @@ class LocalWalFile : public WalFile { std::string wal_path_{}; std::mutex file_mutex_; uint32_t max_docs_wal_flush_{0}; - std::atomic docs_count_{0UL}; + uint64_t docs_count_{0}; WalHeader header_; bool opened_{false}; + bool read_only_{false}; + bool reader_ready_{false}; bool failed_{false}; // Preserve the complete prefix and remove a torn final record before the // next append. Merely reading a WAL must not modify it. diff --git a/src/db/index/storage/wal/wal_file.h b/src/db/index/storage/wal/wal_file.h index 34c2570b3..4b53b8254 100644 --- a/src/db/index/storage/wal/wal_file.h +++ b/src/db/index/storage/wal/wal_file.h @@ -29,6 +29,7 @@ using WalFilePtr = std::shared_ptr; struct WalOptions { uint32_t max_docs_wal_flush{0}; bool create_new{false}; + bool read_only{false}; }; class WalFile { @@ -48,6 +49,7 @@ class WalFile { public: virtual int append(std::string &&data) = 0; + // Prepare before next(), including after append() changes the file position. virtual int prepare_for_read() = 0; // A successful empty optional means EOF or an incomplete final crash record. // Read failures and complete but corrupt records return an error. diff --git a/src/include/zvec/db/doc.h b/src/include/zvec/db/doc.h index 785d9c1d8..2366708dd 100644 --- a/src/include/zvec/db/doc.h +++ b/src/include/zvec/db/doc.h @@ -310,14 +310,9 @@ class ZVEC_API Doc { private: static void serialize_value(std::vector &buffer, const Value &value); - static Value deserialize_value(const uint8_t *&data, uint8_t type); - static Value deserialize_value(const uint8_t *&data); - static void write_to_buffer(std::vector &buffer, const void *src, size_t size); - static void read_from_buffer(const uint8_t *&data, void *dest, size_t size); - struct ValueEqual; private: diff --git a/tests/db/crash_recovery/wal_recovery_test.cc b/tests/db/crash_recovery/wal_recovery_test.cc index ef757408c..05da1bfd1 100644 --- a/tests/db/crash_recovery/wal_recovery_test.cc +++ b/tests/db/crash_recovery/wal_recovery_test.cc @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -558,5 +559,131 @@ TEST_F(WalRecoveryDeathTest, UpsertWalRecordsInsertOrUpdatePredecessor) { EXPECT_EQ(docs[1]->get("text"), "second"); } +class InvalidWalPayloadDeathTest : public WalRecoveryDeathTest, + public ::testing::WithParamInterface {}; + +TEST_P(InvalidWalPayloadDeathTest, FailsBeforeApplyingValidPrefix) { + ASSERT_EXIT(WriteUpsertsAndExit(path_, true, true), + ::testing::ExitedWithCode(0), ""); + std::vector docs; + ASSERT_NO_FATAL_FAILURE(ReadWalDocuments(path_, &docs)); + ASSERT_EQ(docs.size(), 3u); + const int corruption = GetParam(); + if (corruption == 1) { + docs.back()->set_operator(static_cast(99)); + } else if (corruption >= 2) { + docs.back()->set_operator(corruption == 2 ? Operator::UPDATE + : Operator::DELETE); + docs.back()->set_doc_id(std::numeric_limits::max()); + } + const auto wal_path = FindWal(path_); + auto wal = WalFile::Create(wal_path); + ASSERT_EQ(wal->remove(), 0); + WalOptions options; + options.create_new = true; + ASSERT_EQ(wal->open(options), 0); + for (size_t i = 0; i < docs.size(); ++i) { + auto payload = docs[i]->serialize(); + if (corruption == 0 && i + 1 == docs.size()) payload.resize(3); + // Recompute the CRC so validation must inspect the document, not just the + // WAL frame. A valid prefix must remain unapplied even after retrying open. + ASSERT_EQ(wal->append(std::string(payload.begin(), payload.end())), 0); + } + ASSERT_EQ(wal->close(), 0); + const auto idmap_path = FindIdMap(path_); + ASSERT_FALSE(idmap_path.empty()); + { + auto map = IDMap::CreateAndOpen("wal_recovery", idmap_path, false, false); + ASSERT_NE(map, nullptr); + ASSERT_TRUE(map->upsert("target", 0).ok()); + map->remove("broken"); + ASSERT_TRUE(map->flush().ok()); + } + const auto bytes = ReadFileBytes(wal_path); + const auto manifests = ReadManifests(path_); + for (int attempt = 0; attempt < 2; ++attempt) { + auto opened = Collection::Open(path_, CollectionOptions{}); + ASSERT_FALSE(opened.has_value()); + EXPECT_EQ(opened.error().code(), StatusCode::INTERNAL_ERROR); + EXPECT_NE(opened.error().message().find(corruption < 2 + ? "Corrupt WAL document" + : "Invalid WAL predecessor"), + std::string::npos); + EXPECT_EQ(ReadFileBytes(wal_path), bytes); + EXPECT_EQ(ReadManifests(path_), manifests); + auto map = IDMap::CreateAndOpen("wal_recovery", idmap_path, false, true); + ASSERT_NE(map, nullptr); + uint64_t id; + ASSERT_TRUE(map->has("target", &id)); + EXPECT_EQ(id, 0u); + EXPECT_FALSE(map->has("broken")); + } +} + +INSTANTIATE_TEST_SUITE_P(CorruptDocuments, InvalidWalPayloadDeathTest, + ::testing::Values(0, 1, 2, 3)); + +class WalBlockBoundaryDeathTest + : public WalRecoveryDeathTest, + public ::testing::WithParamInterface {}; + +TEST_P(WalBlockBoundaryDeathTest, PendingWriteSurvivesBlockRotation) { + const auto operation = GetParam(); + ASSERT_EXIT( + { + CollectionSchema schema("wal_recovery"); + if (!schema + .add_field(std::make_shared( + "text", DataType::STRING, false)) + .ok()) + std::_Exit(1); + CollectionOptions options; + options.max_buffer_size_ = 1; // The first row fills the block. + auto created = Collection::CreateAndOpen(path_, schema, options); + if (!created.has_value()) std::_Exit(2); + Doc doc; + doc.set_pk("target"); + doc.set("text", "first"); + std::vector first{doc}; + auto inserted = created.value()->insert(first); + if (!inserted.has_value() || !inserted.value().front().ok()) + std::_Exit(3); + if (operation == Operator::INSERT) doc.set_pk("suffix"); + doc.set("text", "second"); + std::vector second{doc}; + auto written = + operation == Operator::INSERT ? created.value()->insert(second) + : operation == Operator::UPDATE ? created.value()->update(second) + : created.value()->upsert(second); + if (!written.has_value() || !written.value().front().ok()) + std::_Exit(4); + std::_Exit(0); + }, + ::testing::ExitedWithCode(0), ""); + ASSERT_FALSE(FindWal(path_).empty()); + for (int reopen = 0; reopen < 2; ++reopen) { + auto opened = Collection::Open(path_, CollectionOptions{}); + ASSERT_TRUE(opened.has_value()) << opened.error().message(); + if (operation == Operator::INSERT) { + auto fetched = opened.value()->fetch({"target", "suffix"}); + ASSERT_TRUE(fetched.has_value()); + ASSERT_NE(fetched.value().at("target"), nullptr); + ASSERT_NE(fetched.value().at("suffix"), nullptr); + EXPECT_EQ(fetched.value().at("target")->get("text"), + "first"); + EXPECT_EQ(fetched.value().at("suffix")->get("text"), + "second"); + EXPECT_EQ(opened.value()->stats().value().doc_count, 2u); + } else { + ASSERT_NO_FATAL_FAILURE(ExpectOnlyTarget(opened.value(), "second")); + } + ASSERT_TRUE(opened.value()->flush().ok()); + } +} + +INSTANTIATE_TEST_SUITE_P(WriteOperations, WalBlockBoundaryDeathTest, + ::testing::Values(Operator::INSERT, Operator::UPDATE, + Operator::UPSERT)); + } // namespace } // namespace zvec diff --git a/tests/db/index/common/doc_test.cc b/tests/db/index/common/doc_test.cc index 32e875651..10d64d2e4 100644 --- a/tests/db/index/common/doc_test.cc +++ b/tests/db/index/common/doc_test.cc @@ -1026,6 +1026,14 @@ TEST_F(DocDetailedTest, SerializeValueCoverage) { auto buffer = doc.serialize(); EXPECT_FALSE(buffer.empty()); + // Exercise every truncated header, scalar, vector and nested string using + // independently allocated buffers, so ASan also catches accidental overreads. + for (size_t size = 0; size < buffer.size(); ++size) { + SCOPED_TRACE(size); + std::vector truncated(buffer.begin(), buffer.begin() + size); + EXPECT_EQ(Doc::deserialize(truncated.data(), truncated.size()), nullptr); + } + auto deserialized_doc = Doc::deserialize(buffer.data(), buffer.size()); EXPECT_NE(deserialized_doc, nullptr); @@ -1674,6 +1682,90 @@ TEST_F(DocDetailedTest, } } +TEST_F(DocDetailedTest, DeserializeRejectsMalformedLengthsAndTags) { + Doc doc; + doc.set_pk("id"); + doc.set("field", "payload"); + const auto valid = doc.serialize(); + const size_t operation = + sizeof(uint32_t) + 2 + sizeof(float) + sizeof(uint64_t); + const size_t fields_count = operation + sizeof(uint32_t); + const size_t field = fields_count + sizeof(uint32_t); + const size_t tag = field + sizeof(uint32_t) + 5; + const auto maximum = std::numeric_limits::max(); + for (size_t offset : {size_t{0}, fields_count, field, tag + 1}) { + SCOPED_TRACE(offset); + auto invalid = valid; + std::memcpy(invalid.data() + offset, &maximum, sizeof(maximum)); + EXPECT_EQ(Doc::deserialize(invalid.data(), invalid.size()), nullptr); + } + auto invalid = valid; + std::memcpy(invalid.data() + operation, &maximum, sizeof(maximum)); + EXPECT_EQ(Doc::deserialize(invalid.data(), invalid.size()), nullptr); + invalid = valid; + invalid[tag] = 255; + EXPECT_EQ(Doc::deserialize(invalid.data(), invalid.size()), nullptr); + invalid = valid; + invalid.push_back(0); + EXPECT_EQ(Doc::deserialize(invalid.data(), invalid.size()), nullptr); + invalid = valid; + const uint32_t two = 2; + std::memcpy(invalid.data() + fields_count, &two, sizeof(two)); + invalid.insert(invalid.end(), valid.begin() + field, valid.end()); + EXPECT_EQ(Doc::deserialize(invalid.data(), invalid.size()), nullptr); + EXPECT_EQ(Doc::deserialize(nullptr, valid.size()), nullptr); + EXPECT_EQ(Doc::deserialize(valid.data(), 1), nullptr); +} + +TEST_F(DocDetailedTest, DeserializeChecksNestedCountsAndBooleanRepresentation) { + const size_t tag = sizeof(uint32_t) + 2 + sizeof(float) + sizeof(uint64_t) + + sizeof(uint32_t) + sizeof(uint32_t) + sizeof(uint32_t) + 1; + const std::vector values{ + std::vector{1, 2}, std::vector{"hello", "world"}, + std::vector{true, false}, + std::pair, std::vector>{{1, 2}, {1.f, 2.f}}, + std::pair, std::vector>{ + {1, 2}, {zvec::float16_t(1.f), zvec::float16_t(2.f)}}}; + for (const auto &value : values) { + Doc doc; + doc.set_pk("id"); + std::visit( + [&](const auto &item) { + if constexpr (std::is_same_v, + std::monostate>) { + doc.set_null("v"); + } else { + doc.set("v", item); + } + }, + value); + auto buffer = doc.serialize(); + ASSERT_NE(Doc::deserialize(buffer.data(), buffer.size()), nullptr); + // A UINT32_MAX element count cannot fit into any of these payloads. + std::fill_n(buffer.begin() + tag + 1, sizeof(uint32_t), uint8_t{0xff}); + EXPECT_EQ(Doc::deserialize(buffer.data(), buffer.size()), nullptr); + } + Doc sparse; + sparse.set_pk("id"); + sparse.set("v", std::pair, std::vector>{ + {1, 2}, {1.f, 2.f}}); + auto buffer = sparse.serialize(); + std::fill_n( + buffer.begin() + tag + 1 + sizeof(uint32_t) + 2 * sizeof(uint32_t), + sizeof(uint32_t), uint8_t{0xff}); + EXPECT_EQ(Doc::deserialize(buffer.data(), buffer.size()), nullptr); + Doc boolean; + boolean.set_pk("id"); + boolean.set("v", true); + buffer = boolean.serialize(); + buffer[tag + 1] = 2; + EXPECT_EQ(Doc::deserialize(buffer.data(), buffer.size()), nullptr); + boolean.set("v", std::vector{true}); + buffer = boolean.serialize(); + buffer[tag + 1 + sizeof(uint32_t)] = 2; + EXPECT_EQ(Doc::deserialize(buffer.data(), buffer.size()), nullptr); +} + TEST_F(DocDetailedTest, DeserializePreservesHistoricalTextWithoutRevalidation) { Doc doc; doc.set_pk(std::string("old\0id", 6)); diff --git a/tests/db/index/storage/wal_file_test.cc b/tests/db/index/storage/wal_file_test.cc index a29e80e3d..43fef65d8 100644 --- a/tests/db/index/storage/wal_file_test.cc +++ b/tests/db/index/storage/wal_file_test.cc @@ -246,6 +246,7 @@ TEST_F(WalFileTest, TestBoundaryCondition) { wal_option.create_new = false; ret = wal_file->open(wal_option); ASSERT_EQ(ret, 0); + ASSERT_EQ(wal_file->prepare_for_read(), 0); uint32_t idx = 0; std::string record = ReadRecord(wal_file); while (!record.empty()) { @@ -430,6 +431,7 @@ TEST_F(WalFileTest, TestFirstErrorCase) { EXPECT_EQ(result.error().code(), StatusCode::INTERNAL_ERROR); EXPECT_NE(result.error().message().find("CRC mismatch"), std::string::npos); EXPECT_NE(wal_file->append("after corruption"), 0); + EXPECT_NE(wal_file->flush(), 0); // close ret = wal_file->close(); ASSERT_EQ(ret, 0); @@ -803,6 +805,59 @@ TEST_F(WalFileTest, ClosedReaderReturnsError) { EXPECT_EQ(result.error().code(), StatusCode::INTERNAL_ERROR); } +TEST_F(WalFileTest, ReadOnlyWalCannotBeModified) { + const std::string path = "./data.wal.readonly"; + auto wal = WalFile::Create(path); + WalOptions options; + options.create_new = true; + ASSERT_EQ(wal->open(options), 0); + ASSERT_EQ(wal->append("prefix"), 0); + ASSERT_EQ(wal->close(), 0); + options.create_new = false; + options.read_only = true; + ASSERT_EQ(wal->open(options), 0); + ASSERT_EQ(wal->prepare_for_read(), 0); + EXPECT_EQ(ReadRecord(wal), "prefix"); + EXPECT_NE(wal->append("suffix"), 0); + EXPECT_NE(wal->flush(), 0); + EXPECT_NE(wal->remove(), 0); + ASSERT_EQ(wal->close(), 0); + EXPECT_NE(wal->remove(), 0); + options.read_only = false; + ASSERT_EQ(wal->open(options), 0); + ASSERT_EQ(wal->prepare_for_read(), 0); + EXPECT_EQ(ReadRecord(wal), "prefix"); + auto end = wal->next(); + ASSERT_TRUE(end.has_value()); + EXPECT_FALSE(end.value().has_value()); + ASSERT_EQ(wal->remove(), 0); +} + +TEST_F(WalFileTest, AppendRequiresReadCursorToBePreparedAgain) { + auto wal = WalFile::Create("./data.wal.cursor"); + WalOptions options; + options.create_new = true; + ASSERT_EQ(wal->open(options), 0); + ASSERT_EQ(wal->append("first"), 0); + ASSERT_EQ(wal->append("second"), 0); + ASSERT_EQ(wal->prepare_for_read(), 0); + EXPECT_EQ(ReadRecord(wal), "first"); + ASSERT_EQ(wal->append("third"), 0); + EXPECT_FALSE(wal->next().has_value()); + ASSERT_EQ(wal->prepare_for_read(), 0); + EXPECT_EQ(ReadRecord(wal), "first"); + EXPECT_EQ(ReadRecord(wal), "second"); + EXPECT_EQ(ReadRecord(wal), "third"); +} + +TEST_F(WalFileTest, RemoveReportsFilesystemFailure) { + const std::string path = "./data.wal.remove"; + ASSERT_TRUE(FileHelper::CreateDirectory(path)); + auto wal = WalFile::Create(path); + EXPECT_NE(wal->remove(), 0); + ASSERT_TRUE(FileHelper::RemoveDirectory(path)); +} + TEST_F(WalFileTest, ZeroLengthRecordIsCorruption) { const std::string path = "./data.wal.zero"; auto wal = WalFile::Create(path); @@ -843,4 +898,4 @@ TEST_F(WalFileTest, TruncatedFileHeaderReturnsError) { #if defined(__GNUC__) || defined(__GNUG__) #pragma GCC diagnostic pop -#endif \ No newline at end of file +#endif