diff --git a/CMakeLists.txt b/CMakeLists.txt index 784b8cad2..7fbf20472 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1766,6 +1766,7 @@ if (BUILD_TESTING) tests/rdb_test.cpp tests/rdb_collection_test.cpp tests/record_location_test.cpp + tests/load_value_awaitable_test.cpp tests/request_gate_test.cpp tests/replication_command_test.cpp tests/replication_group_test.cpp diff --git a/docs/architecture/04-storage-and-recovery.md b/docs/architecture/04-storage-and-recovery.md index 6ba2fdcb6..56a06b7e7 100644 --- a/docs/architecture/04-storage-and-recovery.md +++ b/docs/architecture/04-storage-and-recovery.md @@ -649,7 +649,17 @@ rereads. Transaction-owned reads always use that locked path directly. Tombstones and expired values are invisible; an expired observation can enqueue a bounded active-expiration candidate. A staged location is copied or framed from its write buffer. A disk location is read on its physical block owner into -an aligned lease. +an aligned lease. `LoadValue` and its single-location `LoadValueLocal` share an +awaitable embedded in the caller's task. Local inline reads, including staged +copies, use that task directly. Buffer contention retains the pool's FIFO +handoff and schedules an owner-local continuation under the caller's task class; +it does not require a waiter coroutine. The awaitable owns the disk-read pin +and buffer through completion and validates record identity, checksums, and +population before returning it. Grouped reads enter their graph reader directly; +cross-owner, external, and relocation flows use lazy helper tasks. A stale +logical read resolves the current same-mutation location before resuming its +caller, while a single-location read returns the stale status to its enclosing +representation-specific retry loop. Worker-local MGET uses `BatchGetLocked` after the command has acquired its key locks. It classifies index entries in one coroutine and issues ordinary-size @@ -668,9 +678,13 @@ taking the reader's key lock; a stale physical read aborts and follows the current same-sequence relocation, while a true logical mutation becomes a normal miss. -External values read and validate every extent on its current owner in bounded -parallel waves, then assemble the logical value. Every read in a wave finishes -before its shared output buffer can be released, including on error. A move-only +External values read and validate every extent on its current owner and +assemble the logical value in manifest order. Whole-extent and slice reads +share a physical-read awaitable that retains its pin across buffer waiting, +I/O, and validation. Slices validate the entire extent before copying any +requested bytes. Local extent routing and physical reads need no child task; +cross-owner dispatch retains the runtime's task bridge. The output buffer +outlives every extent read, including on error. A move-only `ReadBufferLease` may cross to the connection worker; destruction returns a registered slot to its storage owner. diff --git a/include/lavik/storage/buffer_pool.h b/include/lavik/storage/buffer_pool.h index 759d2acb4..44bcd9f10 100644 --- a/include/lavik/storage/buffer_pool.h +++ b/include/lavik/storage/buffer_pool.h @@ -32,6 +32,7 @@ namespace bycorf { class CrossCore; class Worker; +struct RemoteWork; } // namespace bycorf namespace lavik::storage { @@ -208,12 +209,20 @@ class RegisteredBufferPool { bool await_suspend(std::coroutine_handle<> awaiting); absl::StatusOr await_resume(); + // Queue an owner-local continuation when a slot is handed off. This uses + // the worker's scheduled work queue (including its task-class budget), + // never runs inline in Release(), and sends no cross-worker reply. The + // caller supplies run_fn_/task_class_ and keeps both objects alive until + // that callback consumes the assigned lease through await_resume(). + bool SuspendWith(bycorf::RemoteWork& continuation); + private: friend class RegisteredBufferPool; RegisteredBufferPool* pool_ = nullptr; std::size_t minimum_payload_bytes_ = 0; std::coroutine_handle<> awaiting_{}; + bycorf::RemoteWork* continuation_ = nullptr; std::uint16_t assigned_buffer_id_ = 0; }; diff --git a/include/lavik/storage/engine.h b/include/lavik/storage/engine.h index 2feb57526..47deb7580 100644 --- a/include/lavik/storage/engine.h +++ b/include/lavik/storage/engine.h @@ -1944,6 +1944,7 @@ class StorageEngine { private: friend class ExpirationAuthorityTestPeer; friend class TombRaiderTestPeer; + friend class LoadValueAwaitableTestPeer; class Impl; std::unique_ptr impl_; }; diff --git a/src/storage/buffer_pool.cpp b/src/storage/buffer_pool.cpp index 9a9ddf292..8d5f34ef1 100644 --- a/src/storage/buffer_pool.cpp +++ b/src/storage/buffer_pool.cpp @@ -468,16 +468,22 @@ bool RegisteredBufferPool::AcquireReadAwaiter::await_ready() const noexcept { bool RegisteredBufferPool::AcquireReadAwaiter::await_suspend( std::coroutine_handle<> awaiting) { - if (pool_ == nullptr || !pool_->initialized() || - minimum_payload_bytes_ > pool_->options_.read_payload_bytes() || - !pool_->free_read_buffers_.empty()) { - return false; - } + if (await_ready()) return false; awaiting_ = awaiting; pool_->read_waiters_.push_back(this); return true; } +bool RegisteredBufferPool::AcquireReadAwaiter::SuspendWith( + bycorf::RemoteWork& continuation) { + if (await_ready()) return false; + continuation_ = &continuation; + continuation.origin_ = pool_->owner_worker_; + continuation.reply_deferred_ = true; + pool_->read_waiters_.push_back(this); + return true; +} + absl::StatusOr RegisteredBufferPool::AcquireReadAwaiter::await_resume() { if (pool_ == nullptr || !pool_->initialized()) { @@ -717,7 +723,9 @@ void RegisteredBufferPool::ReleaseLocal(std::uint16_t buffer_id) noexcept { AcquireReadAwaiter* waiter = read_waiters_.front(); read_waiters_.pop_front(); waiter->assigned_buffer_id_ = buffer_id; - if (worker_ != nullptr) { + if (waiter->continuation_ != nullptr) { + bycorf::PostRequest(cross_core_, owner_worker_, waiter->continuation_); + } else if (worker_ != nullptr) { worker_->Enqueue(waiter->awaiting_); } return; diff --git a/src/storage/engine/impl.h b/src/storage/engine/impl.h index 17dc6ed40..83f39e248 100644 --- a/src/storage/engine/impl.h +++ b/src/storage/engine/impl.h @@ -3257,6 +3257,7 @@ class StorageEngine::Impl { private: friend class ExpirationAuthorityTestPeer; friend class TombRaiderTestPeer; + friend class LoadValueAwaitableTestPeer; struct DurableSystemState { std::uint64_t generation_ = 0; @@ -3342,6 +3343,118 @@ class StorageEngine::Impl { } }; + // The awaitable lives in the caller's frame. Ordinary local reads suspend + // that caller directly, including through pool backpressure. Only multi-step + // logical reads need a lazy helper Task. Local mode returns stale-location + // errors to its caller, while logical mode resolves same-mutation relocations + // before returning. + class LoadValueAwaitable final : public bycorf::OneShotIoAwaitable, + private bycorf::RemoteWork { + public: + using Result = absl::StatusOr; + LoadValueAwaitable(Impl& engine, WorkerStore& store, + WorkerStore::PartitionStore& partition, + std::uint8_t db_id, std::string_view key, + const Digest& digest, RecordLocation location, + ExtentManifest extents, ReadLatencyTrace* trace, + GroupedObject::Handle grouped_snapshot); + LoadValueAwaitable(Impl& engine, WorkerStore& store, std::uint8_t db_id, + std::string_view key, RecordLocation location, + std::uint64_t replication_epoch, ReadLatencyTrace* trace, + std::optional expected_db_epoch); + ~LoadValueAwaitable(); + + bool await_ready() const noexcept { return ready_ && !slow_.valid(); } + std::coroutine_handle<> await_suspend(std::coroutine_handle<> awaiting); + Result await_resume(); + void Complete(bycorf::Worker& worker, int result, unsigned flags) override; + + private: + friend class LoadValueAwaitableTestPeer; + struct Payload { + std::uint64_t txid_; + std::uint32_t offset_; + std::uint32_t bytes_; + std::uint32_t checksum_; + }; + + void PrepareLocal(); + void FinishBufferWait(); + static void BufferReady(bycorf::RemoteWork* work); + std::coroutine_handle<> SuspendRead(std::coroutine_handle<> awaiting); + absl::StatusOr ParseRecord(const std::byte* data) const; + Result FinishRecord(const std::byte* data, Payload payload, bool staged); + bool PopulationChanged() const; + void Unpin(); + void Retry(absl::Status failure); + + Impl& engine_; + WorkerStore& store_; + WorkerStore::PartitionStore* partition_ = nullptr; + std::string_view key_; + Digest digest_{}; + RecordLocation location_; + ReadLatencyTrace* trace_; + std::uint64_t replication_epoch_; + std::optional expected_db_epoch_; + std::uint64_t aligned_offset_ = 0; + std::size_t file_id_ = 0; + std::size_t record_headroom_ = 0; + std::size_t read_bytes_ = 0; + std::uint8_t db_id_; + BlockState* pinned_ = nullptr; + std::optional payload_; + std::optional ready_; + // Runtime shutdown drains I/O before destroying suspended frames. The + // waiter retains its assigned slot until the scheduled continuation runs. + ReadBufferLease lease_; + std::optional buffer_wait_; + Task slow_; + }; + + // One physical extent read, including buffer waiting and slice validation, + // embedded in the caller. Routing adds a runtime task bridge only when the + // block belongs to another worker; local reads need no child coroutine. + class ReadExtentAwaitable final : public bycorf::OneShotIoAwaitable, + private bycorf::RemoteWork { + public: + ReadExtentAwaitable(Impl& engine, WorkerStore& store, ExtentRef ref, + std::uint32_t extent_index, std::size_t source_offset, + std::span destination, bool route = false); + ~ReadExtentAwaitable(); + bool await_ready() const noexcept { return ready_.has_value(); } + bool await_suspend(std::coroutine_handle<> awaiting); + absl::Status await_resume(); + + private: + friend class LoadValueAwaitableTestPeer; + struct RemoteRead { + Impl* engine_; + unsigned owner_; + ExtentRef ref_; + std::uint32_t extent_index_; + std::size_t source_offset_; + std::span destination_; + Task operator()(); + }; + void FinishBufferWait(); + static void BufferReady(bycorf::RemoteWork* work); + bool SubmitRead(std::coroutine_handle<> awaiting); + + Impl& engine_; + WorkerStore& store_; + ExtentRef ref_; + std::uint32_t extent_index_; + std::size_t source_offset_; + std::span destination_; + std::size_t read_bytes_ = 0; + BlockState* pinned_ = nullptr; + std::optional ready_; + ReadBufferLease lease_; + std::optional buffer_wait_; + std::optional> remote_; + }; + // Owns the page lease while a grouped reader inspects its checked envelope. // Views into loaded_.value() must not outlive this object. struct LoadedHashGroupPayload { @@ -3726,13 +3839,21 @@ class StorageEngine::Impl { WorkerStore& store, RecordIndex& index, const Digest& digest, std::string_view key, const RecordLocation& verified_location); - Task> LoadValue( + LoadValueAwaitable LoadValue( WorkerStore& key_store, WorkerStore::PartitionStore& partition, std::uint8_t db_id, std::string_view key, const Digest& digest, RecordLocation location, ExtentManifest extents, ReadLatencyTrace* trace = nullptr, GroupedObject::Handle grouped_snapshot = nullptr); + // Used only for multi-step reads or an aborted direct read. A supplied + // failure enters relocation resolution without reading the stale bytes again. + Task> LoadValueSlow( + WorkerStore& key_store, WorkerStore::PartitionStore& partition, + std::uint8_t db_id, std::string_view key, Digest digest, + RecordLocation location, ExtentManifest extents, ReadLatencyTrace* trace, + std::optional failure = std::nullopt); + // The optional snapshot view must already own physical pins. Ordinary // reads instead retry GC relocation against the same logical root version. Task> LoadHashGroupPayload( @@ -3752,26 +3873,22 @@ class StorageEngine::Impl { std::uint8_t db_id, std::string_view key, const Digest& digest, RecordLocation location, GroupedObject::Handle snapshot); - // Reads one extent block's payload into `destination`. Runs on the worker - // that owns that block, which is not necessarily the one holding the - // manifest, so everything it needs is passed by value. - Task ReadExtentInto(WorkerStore& store, ExtentRef ref, - std::uint32_t extent_index, - std::byte* destination); - Task ReadExtentOnOwner(ExtentRef ref, - std::uint32_t extent_index, - std::byte* destination); - Task ReadExtentSlice(WorkerStore& store, ExtentRef ref, - std::uint32_t extent_index, - std::size_t source_offset, - std::span destination); + // Reads an entire extent on its physical owner into caller-owned storage. + // Local routing returns the physical awaitable without a wrapper task. + ReadExtentAwaitable ReadExtentOnOwner(ExtentRef ref, + std::uint32_t extent_index, + std::byte* destination); + ReadExtentAwaitable ReadExtentSlice(WorkerStore& store, ExtentRef ref, + std::uint32_t extent_index, + std::size_t source_offset, + std::span destination); Task> LoadExternalValueLocal( WorkerStore& store, const RecordLocation& location, ExtentManifest extents, ReadLatencyTrace* trace, bool grouped_payload = false); - Task> LoadValueLocal( + LoadValueAwaitable LoadValueLocal( WorkerStore& store, std::uint8_t db_id, std::string_view key, RecordLocation location, std::uint64_t replication_epoch, ReadLatencyTrace* trace = nullptr, diff --git a/src/storage/engine/read.cpp b/src/storage/engine/read.cpp index b987a39b1..6d9401122 100644 --- a/src/storage/engine/read.cpp +++ b/src/storage/engine/read.cpp @@ -879,17 +879,14 @@ absl::StatusOr StorageEngine::Impl::EncodeDiskValue( } Task> -StorageEngine::Impl::LoadValue(WorkerStore& key_store, - WorkerStore::PartitionStore& partition, - std::uint8_t db_id, std::string_view key, - const Digest& digest, RecordLocation location, - ExtentManifest extents, ReadLatencyTrace* trace, - GroupedObject::Handle grouped_snapshot) { - if (location.grouped()) { - co_return co_await LoadGroupedValue(key_store, partition, db_id, key, - digest, location, - std::move(grouped_snapshot)); - } +StorageEngine::Impl::LoadValueSlow(WorkerStore& key_store, + WorkerStore::PartitionStore& partition, + std::uint8_t db_id, std::string_view key, + Digest digest, RecordLocation location, + ExtentManifest extents, + ReadLatencyTrace* trace, + std::optional failure) { + assert(!location.grouped()); // The caller already resolved the key's partition for the index lookup. // Reuse it across retries instead of recomputing the Redis slot. const std::uint64_t replication_epoch = partition.replication_epoch_; @@ -901,9 +898,8 @@ StorageEngine::Impl::LoadValue(WorkerStore& key_store, : std::nullopt; while (true) { assert(location.block_owner() < worker_count_); - // Every branch below assigns `loaded` before it is observed. Keep the - // required placeholder message-free so the common local GET path does not - // construct and then immediately discard an error payload. + // Every branch assigns `loaded` before it is observed; the placeholder + // does not need an allocated error message. absl::StatusOr loaded; // An external value's manifest is already decoded in this index entry, @@ -911,7 +907,10 @@ StorageEngine::Impl::LoadValue(WorkerStore& key_store, // and let each extent go straight to its block's owner, rather than // hopping to the record's owner first and having it acquire the output // buffer from its pool and hand the lease back across workers. - if (location.external()) { + if (failure) { + loaded = std::move(*failure); + failure.reset(); + } else if (location.external()) { loaded = co_await LoadExternalValueLocal(key_store, location, std::move(extents), trace); } else if (location.block_owner() == key_store.worker_->id()) { @@ -973,131 +972,155 @@ StorageEngine::Impl::LoadValue(WorkerStore& key_store, } } -Task StorageEngine::Impl::ReadExtentInto( - WorkerStore& store, ExtentRef ref, std::uint32_t extent_index, - std::byte* destination) { - BlockState* state = FindBlockState(store, ref.block_id_); - if (state == nullptr || !state->allocated_ || state->freeing_ || - state->kind_ != BlockKind::kPayloadExtent || - state->allocation_epoch_ != ref.allocation_epoch_) { - co_return absl::Status(absl::StatusCode::kAborted, - "stale or missing external extent"); - } - ++state->pins_; - struct ExtentPin { - BlockState* state_; - ~ExtentPin() { --state_->pins_; } - } pin{state}; - const std::size_t read_bytes = - AlignDirect(kBlockHeaderBytes + ref.payload_bytes_); - auto temp_acquired = co_await store.buffers_.AcquireReadBuffer(read_bytes); - if (!temp_acquired.ok()) { - co_return temp_acquired.status(); - } - ReadBufferLease temp = std::move(*temp_acquired); - FixedBuffer io = temp.io_buffer(); - io.size_ = read_bytes; - const auto [file_id, block_offset] = FileOffset(ref.block_id_); - auto read = co_await ReadStorageBuffer(*store.worker_, store.files_[file_id], - io, temp.registered(), block_offset); - if (!read.ok()) { - co_return read.status(); - } - if (*read != read_bytes) { - co_return absl::Status(absl::StatusCode::kInternal, - "short extent block read"); - } - BlockHeader header{}; - if (!DecodeBlockHeaderPages(std::span( - io.data_, kBlockHeaderBytes), - &header) || - header.kind_ != BlockKind::kPayloadExtent || - header.block_id_ != ref.block_id_ || - header.allocation_epoch_ != ref.allocation_epoch_ || - header.extent_index_ != extent_index || - header.extent_payload_bytes_ != ref.payload_bytes_ || - header.extent_payload_checksum_ != ref.payload_checksum_) { - co_return absl::Status(absl::StatusCode::kInternal, - "extent header does not match manifest"); - } - const auto payload = std::span(io.data_ + kBlockHeaderBytes, - ref.payload_bytes_); - if (Crc32c(payload) != ref.payload_checksum_) { - co_return absl::Status(absl::StatusCode::kInternal, - "extent payload checksum mismatch"); - } - std::memcpy(destination, payload.data(), payload.size()); - co_return absl::OkStatus(); -} - -Task StorageEngine::Impl::ReadExtentOnOwner( +StorageEngine::Impl::ReadExtentAwaitable StorageEngine::Impl::ReadExtentOnOwner( ExtentRef ref, std::uint32_t extent_index, std::byte* destination) { - const unsigned owner = BlockOwner(ref.block_id_); - // The block may be recycled after manifest validation. Recheck its owner - // before routing, and let ReadExtentInto validate the allocation epoch. - if (owner >= worker_count_) - co_return absl::AbortedError("stale or missing external extent"); - if (owner == bycorf::ThisWorker().id_) - co_return co_await ReadExtentInto(*stores_[owner], ref, extent_index, - destination); - co_return co_await bycorf::SubmitTaskTo( - owner, [this, owner, ref, extent_index, destination]() { - return ReadExtentInto(*stores_[owner], ref, extent_index, destination); - }); + return ReadExtentAwaitable(*this, CurrentStore(), ref, extent_index, 0, + std::span(destination, ref.payload_bytes_), true); } -Task StorageEngine::Impl::ReadExtentSlice( +StorageEngine::Impl::ReadExtentAwaitable StorageEngine::Impl::ReadExtentSlice( WorkerStore& store, ExtentRef ref, std::uint32_t extent_index, std::size_t source_offset, std::span destination) { + return ReadExtentAwaitable(*this, store, ref, extent_index, source_offset, + destination); +} + +Task +StorageEngine::Impl::ReadExtentAwaitable::RemoteRead::operator()() { + // SubmitTaskTo requires a Task on the physical owner. The local awaitable + // below supplies the I/O and pool continuation without another loader frame. + co_return co_await engine_->ReadExtentSlice(*engine_->stores_[owner_], ref_, + extent_index_, source_offset_, + destination_); +} + +StorageEngine::Impl::ReadExtentAwaitable::ReadExtentAwaitable( + Impl& engine, WorkerStore& store, ExtentRef ref, std::uint32_t extent_index, + std::size_t source_offset, std::span destination, bool route) + : engine_(engine), + store_(store), + ref_(ref), + extent_index_(extent_index), + source_offset_(source_offset), + destination_(destination) { if (source_offset > ref.payload_bytes_ || destination.size() > ref.payload_bytes_ - source_offset) { - co_return absl::InvalidArgumentError("extent slice is out of range"); + ready_ = absl::InvalidArgumentError("extent slice is out of range"); + return; + } + if (route) { + const unsigned owner = engine.BlockOwner(ref.block_id_); + if (owner >= engine.worker_count_) { + ready_ = absl::AbortedError("stale or missing external extent"); + return; + } + if (owner != store.worker_->id()) { + remote_.emplace(bycorf::SubmitTaskTo( + owner, RemoteRead{&engine, owner, ref, extent_index, source_offset, + destination})); + return; + } } - BlockState* state = FindBlockState(store, ref.block_id_); - if (state == nullptr || !state->allocated_ || state->freeing_ || + // Recheck the physical owner and incarnation after routing. Retain the + // extent pin across pool waiting as well as I/O, matching manifest reads. + auto* state = engine.FindBlockState(store, ref.block_id_); + if (!state || !state->allocated_ || state->freeing_ || state->kind_ != BlockKind::kPayloadExtent || state->allocation_epoch_ != ref.allocation_epoch_) { - co_return absl::AbortedError("stale or missing external extent"); + ready_ = absl::AbortedError("stale or missing external extent"); + return; } ++state->pins_; - struct ExtentPin { - BlockState* state_; - ~ExtentPin() { --state_->pins_; } - } pin{state}; - const std::size_t read_bytes = - AlignDirect(kBlockHeaderBytes + ref.payload_bytes_); - auto acquired = co_await store.buffers_.AcquireReadBuffer(read_bytes); - if (!acquired.ok()) co_return acquired.status(); - ReadBufferLease buffer = std::move(*acquired); - FixedBuffer io = buffer.io_buffer(); - io.size_ = read_bytes; - const auto [file_id, block_offset] = FileOffset(ref.block_id_); - auto read = co_await ReadStorageBuffer(*store.worker_, store.files_[file_id], - io, buffer.registered(), block_offset); - if (!read.ok()) co_return read.status(); - if (*read != read_bytes) { - co_return absl::InternalError("short extent block read"); + pinned_ = state; + read_bytes_ = AlignDirect(kBlockHeaderBytes + ref.payload_bytes_); + buffer_wait_.emplace(&store.buffers_, read_bytes_); + if (buffer_wait_->await_ready()) FinishBufferWait(); +} + +StorageEngine::Impl::ReadExtentAwaitable::~ReadExtentAwaitable() { + if (pinned_) --pinned_->pins_; +} + +void StorageEngine::Impl::ReadExtentAwaitable::FinishBufferWait() { + auto acquired = buffer_wait_->await_resume(); + buffer_wait_.reset(); + if (!acquired.ok()) + ready_ = acquired.status(); + else + lease_ = std::move(*acquired); +} + +void StorageEngine::Impl::ReadExtentAwaitable::BufferReady( + bycorf::RemoteWork* work) { + auto& read = *static_cast(work); + read.FinishBufferWait(); + if (read.ready_ || !read.SubmitRead(read.awaiting_)) + read.store_.worker_->Enqueue(read.awaiting_); +} + +bool StorageEngine::Impl::ReadExtentAwaitable::await_suspend( + std::coroutine_handle<> awaiting) { + if (remote_) { + remote_->await_suspend(awaiting); + return true; } + if (buffer_wait_) { + awaiting_ = awaiting; + run_fn_ = &BufferReady; + task_class_ = store_.worker_->IsBackground(awaiting) + ? bycorf::TaskClass::kBackground + : bycorf::CurrentTaskClass(); + if (buffer_wait_->SuspendWith(*this)) return true; + FinishBufferWait(); + if (ready_) return false; + } + return SubmitRead(awaiting); +} + +bool StorageEngine::Impl::ReadExtentAwaitable::SubmitRead( + std::coroutine_handle<> awaiting) { + auto buffer = lease_.io_buffer(); + buffer.size_ = read_bytes_; + const auto [file_id, offset] = engine_.FileOffset(ref_.block_id_); + return Suspend( + awaiting, lease_.registered() + ? store_.worker_->SubmitReadFixed(store_.files_[file_id], + buffer, offset, this) + : store_.worker_->SubmitRead( + store_.files_[file_id], + std::span(buffer.data_, buffer.size_), offset, this)); +} + +absl::Status StorageEngine::Impl::ReadExtentAwaitable::await_resume() { + if (remote_) return remote_->await_resume(); + if (ready_) return std::move(*ready_); + auto read = + Resume(lease_.registered() ? "fixed-buffer read failed" : "read failed"); + if (!read.ok()) return read.status(); + store_.worker_->RecordStorageReadCompletion(*read); + if (static_cast(*read) != read_bytes_) + return absl::InternalError("short extent block read"); + const auto buffer = lease_.io_buffer(); BlockHeader header{}; if (!DecodeBlockHeaderPages(std::span( - io.data_, kBlockHeaderBytes), + buffer.data_, kBlockHeaderBytes), &header) || header.kind_ != BlockKind::kPayloadExtent || - header.block_id_ != ref.block_id_ || - header.allocation_epoch_ != ref.allocation_epoch_ || - header.extent_index_ != extent_index || - header.extent_payload_bytes_ != ref.payload_bytes_ || - header.extent_payload_checksum_ != ref.payload_checksum_) { - co_return absl::InternalError("extent header does not match manifest"); - } - const auto payload = std::span(io.data_ + kBlockHeaderBytes, - ref.payload_bytes_); - if (Crc32c(payload) != ref.payload_checksum_) { - co_return absl::InternalError("extent payload checksum mismatch"); - } - std::memcpy(destination.data(), payload.data() + source_offset, - destination.size()); - co_return absl::OkStatus(); + header.block_id_ != ref_.block_id_ || + header.allocation_epoch_ != ref_.allocation_epoch_ || + header.extent_index_ != extent_index_ || + header.extent_payload_bytes_ != ref_.payload_bytes_ || + header.extent_payload_checksum_ != ref_.payload_checksum_) + return absl::InternalError("extent header does not match manifest"); + const auto payload = std::span( + buffer.data_ + kBlockHeaderBytes, ref_.payload_bytes_); + // Even a partial read validates the entire extent before publishing bytes. + // This runs in the resumed task's budget, never on the I/O completion poller. + if (Crc32c(payload) != ref_.payload_checksum_) + return absl::InternalError("extent payload checksum mismatch"); + std::memcpy(destination_.data(), payload.data() + source_offset_, + destination_.size()); + return absl::OkStatus(); } Task> StorageEngine::Impl::LoadOutOfIndexKey( @@ -1299,223 +1322,330 @@ StorageEngine::Impl::LoadExternalValueLocal(WorkerStore& store, co_return LoadedValue{std::move(output), value_offset, output_offset}; } -Task> -StorageEngine::Impl::LoadValueLocal( +StorageEngine::Impl::LoadValueAwaitable StorageEngine::Impl::LoadValue( + WorkerStore& store, WorkerStore::PartitionStore& partition, + std::uint8_t db_id, std::string_view key, const Digest& digest, + RecordLocation location, ExtentManifest extents, ReadLatencyTrace* trace, + GroupedObject::Handle grouped_snapshot) { + return LoadValueAwaitable(*this, store, partition, db_id, key, digest, + location, std::move(extents), trace, + std::move(grouped_snapshot)); +} + +StorageEngine::Impl::LoadValueAwaitable StorageEngine::Impl::LoadValueLocal( WorkerStore& store, std::uint8_t db_id, std::string_view key, RecordLocation location, std::uint64_t replication_epoch, ReadLatencyTrace* trace, std::optional expected_db_epoch) { - LAVIK_FAULT_INJECT(if (LAVIK_FAULT_MATCHES("LAVIK_FAIL_VALUE_READ_KEY", - key)) co_return absl:: - InternalError("injected value payload read failure");); - if (location.external()) { - co_return absl::Status(absl::StatusCode::kInternal, - "external value was dispatched as inline"); - } - // TODO: Coalesce concurrent reads of the same aligned disk page. Key an - // in-flight table by - // (file_id, aligned offset, aligned length), submit one read, and fan the - // decoded result out to all waiting coroutines. In-flight operations must - // retain values/leases, never flat_hash_map iterators or element pointers. - const auto [file_id, block_offset] = FileOffset(location.block_id()); - const std::uint64_t absolute_offset = block_offset + location.record_offset(); - const std::uint64_t direct_io_mask = - static_cast(direct_io_alignment_ - 1); - const std::uint64_t aligned_offset = absolute_offset & ~direct_io_mask; - const std::size_t record_headroom = - static_cast(absolute_offset - aligned_offset); - const std::size_t record_span = record_headroom + location.total_disk_bytes(); - const std::size_t read_bytes = - (record_span + direct_io_alignment_ - 1) & ~direct_io_mask; - - // Acquire the output buffer before resolving any block state. This is the - // only suspension the staged-copy path would otherwise have, and it used - // to sit between reading the staging pointer and using it, so that path - // needed a pin to stop a flush from handing the staging buffer back. With - // the acquisition hoisted, the staged copy runs straight through on a - // worker that cannot preempt it, and only the disk read below pins. The - // disk sizing covers the staged copy too: read_bytes is at least - // total_disk_bytes, which is header plus payload. - if (trace != nullptr) { - trace->buffer_acquire_start_ns_ = ReadTraceNowNanos(); - } - auto acquired = co_await store.buffers_.AcquireReadBuffer(read_bytes); - if (!acquired.ok()) { - co_return acquired.status(); - } - ReadBufferLease lease = std::move(*acquired); - if (trace != nullptr) { - trace->buffer_acquired_ns_ = ReadTraceNowNanos(); - trace->heap_read_buffer_ = lease.buffer_id() == 0; - } + return LoadValueAwaitable(*this, store, db_id, key, location, + replication_epoch, trace, expected_db_epoch); +} - BlockState* state = FindBlockState(store, location.block_id()); - if (state == nullptr || !state->allocated_ || state->freeing_ || - state->allocation_epoch_ != location.allocation_epoch()) { - co_return absl::Status(absl::StatusCode::kAborted, - "stale index block epoch"); +StorageEngine::Impl::LoadValueAwaitable::LoadValueAwaitable( + Impl& engine, WorkerStore& store, WorkerStore::PartitionStore& partition, + std::uint8_t db_id, std::string_view key, const Digest& digest, + RecordLocation location, ExtentManifest extents, ReadLatencyTrace* trace, + GroupedObject::Handle grouped_snapshot) + : engine_(engine), + store_(store), + // Grouped snapshots have their own population/pin protocol. + partition_(location.grouped() ? nullptr : &partition), + key_(key), + digest_(digest), + location_(location), + trace_(trace), + replication_epoch_(partition.replication_epoch_), + expected_db_epoch_(engine.EffectiveRecordDbEpoch(partition, db_id)), + db_id_(db_id) { + if (location.grouped()) { + // The grouped task owns its graph/version protocol. Keep its borrowed + // digest in this awaitable, not in the constructor's argument lifetime. + slow_ = engine.LoadGroupedValue(store, partition, db_id, key, digest_, + location, std::move(grouped_snapshot)); + } else if (location.external() || + location.block_owner() != store.worker_->id()) { + slow_ = engine.LoadValueSlow(store, partition, db_id, key, digest, location, + std::move(extents), trace); + } else { + PrepareLocal(); + if (ready_ && absl::IsAborted(ready_->status())) Retry(ready_->status()); } +} - // Only records appended since the last flush live in a staging buffer, so - // on a read-mostly workload this branch is rare. Keep the disk read on the - // straight-line path. - if (location.in_memory() && state->in_memory_) [[unlikely]] { - auto in_mem_buffer = StagingBufferFor(store, *state); - if (!in_mem_buffer.data_ || in_mem_buffer.size_ == 0 || - location.record_offset() + location.total_disk_bytes() > - in_mem_buffer.size_) { - co_return absl::Status(absl::StatusCode::kInternal, - "invalid in-memory location"); - } - if (trace != nullptr) { - trace->io_submit_ns_ = trace->buffer_acquired_ns_; - trace->io_complete_ns_ = trace->buffer_acquired_ns_; - } - FixedBuffer io = lease.io_buffer(); - if (location.external()) { - co_return absl::Status(absl::StatusCode::kInternal, - "external value requires extent loading"); - } - const std::byte* record_bytes = - in_mem_buffer.data_ + location.record_offset(); - RecordHeader record{}; - std::string_view disk_key; - if (!DecodeRecordHeader(std::span( - record_bytes, location.total_disk_bytes()), - &record, &disk_key) || - record.db_id_ != db_id || (!record.key_indirect_ && disk_key != key) || - record.kind_ != RecordKind::kValue || - record.db_epoch_ != expected_db_epoch.value_or(DbEpoch(db_id)) || - record.mutation_sequence_ != location.mutation_sequence_ || - record.replication_epoch_ != replication_epoch || - record.allocation_epoch_ != location.allocation_epoch() || - record.expire_at_ms_ != location.expire_at_ms_ || - record.value_type_ != location.value_type() || - record.external_ != location.external() || - record.key_indirect_ != location.key_indirect() || - location.logical_size_ != record.logical_size_ || - location.total_disk_bytes() != record.total_disk_bytes_) { - co_return absl::Status(absl::StatusCode::kAborted, - "record does not match in-memory location"); - } - const std::byte* payload_data = record_bytes + record.header_bytes_; - if (record.key_indirect_ && record.key_bytes_ != key.size()) { - co_return absl::Status(absl::StatusCode::kAborted, - "record key does not match location"); +StorageEngine::Impl::LoadValueAwaitable::LoadValueAwaitable( + Impl& engine, WorkerStore& store, std::uint8_t db_id, std::string_view key, + RecordLocation location, std::uint64_t replication_epoch, + ReadLatencyTrace* trace, std::optional expected_db_epoch) + : engine_(engine), + store_(store), + key_(key), + location_(location), + trace_(trace), + replication_epoch_(replication_epoch), + expected_db_epoch_(expected_db_epoch), + db_id_(db_id) { + PrepareLocal(); +} + +StorageEngine::Impl::LoadValueAwaitable::~LoadValueAwaitable() { Unpin(); } + +void StorageEngine::Impl::LoadValueAwaitable::Unpin() { + auto* state = std::exchange(pinned_, nullptr); + if (!state) return; + --state->pins_; + if (state->pins_ == 0 && state->release_pending_) + ReleaseStagingBuffer(store_, *state); +} + +void StorageEngine::Impl::LoadValueAwaitable::PrepareLocal() { + LAVIK_FAULT_INJECT(if (LAVIK_FAULT_MATCHES("LAVIK_FAIL_VALUE_READ_KEY", + key_)) { + ready_.emplace(absl::InternalError("injected value payload read failure")); + return; + }); + if (location_.external()) { + ready_.emplace( + absl::InternalError("external value was dispatched as inline")); + return; + } + const auto [file_id, block_offset] = engine_.FileOffset(location_.block_id()); + file_id_ = file_id; + const auto absolute = block_offset + location_.record_offset(); + const auto mask = + static_cast(engine_.direct_io_alignment_ - 1); + aligned_offset_ = absolute & ~mask; + record_headroom_ = absolute - aligned_offset_; + read_bytes_ = + (record_headroom_ + location_.total_disk_bytes() + mask) & ~mask; + if (!lease_.valid()) { + if (trace_) trace_->buffer_acquire_start_ns_ = ReadTraceNowNanos(); + auto acquire = store_.buffers_.AcquireReadBuffer(read_bytes_); + if (!acquire.await_ready()) { + buffer_wait_.emplace(&store_.buffers_, read_bytes_); + return; } - const std::size_t value_bytes = record.payload_bytes_; - if (record.value_type_ == ValueType::kString && !record.grouped_ && - !record.auxiliary_group_ && value_bytes != record.logical_size_) { - co_return absl::Status(absl::StatusCode::kInternal, - "inline value length does not match metadata"); + auto acquired = acquire.await_resume(); + if (!acquired.ok()) { + ready_.emplace(acquired.status()); + return; } - std::memcpy(io.data_, payload_data, value_bytes); - if (Crc32c(std::span( - payload_data, record.payload_bytes_)) != record.payload_checksum_) { - co_return absl::Status(absl::StatusCode::kInternal, - "record value checksum mismatch"); + lease_ = std::move(*acquired); + } + if (trace_) { + trace_->buffer_acquired_ns_ = ReadTraceNowNanos(); + trace_->heap_read_buffer_ = lease_.buffer_id() == 0; + } + // Resolve staging memory only after acquiring the buffer: waiting can let + // flush/defrag retire it. The staged copy itself never suspends, so only + // disk I/O needs a physical pin. + auto* state = engine_.FindBlockState(store_, location_.block_id()); + if (!state || !state->allocated_ || state->freeing_ || + state->allocation_epoch_ != location_.allocation_epoch()) { + ready_.emplace(absl::AbortedError("stale index block epoch")); + return; + } + if (location_.in_memory() && state->in_memory_) [[unlikely]] { + auto buffer = engine_.StagingBufferFor(store_, *state); + if (!buffer.data_ || + location_.record_offset() + location_.total_disk_bytes() > + buffer.size_) { + ready_.emplace(absl::InternalError("invalid in-memory location")); + return; } - if (trace != nullptr) { - trace->decode_done_ns_ = ReadTraceNowNanos(); + if (trace_) { + trace_->io_submit_ns_ = trace_->buffer_acquired_ns_; + trace_->io_complete_ns_ = trace_->buffer_acquired_ns_; } - const std::size_t value_offset = - static_cast(io.data_ - lease.bytes().data()); - co_return LoadedValue{std::move(lease), value_offset, value_bytes, - record.txid_}; + const auto* data = buffer.data_ + location_.record_offset(); + auto payload = ParseRecord(data); + if (!payload.ok()) + ready_.emplace(payload.status()); + else + ready_.emplace(FinishRecord(data, *payload, true)); + return; + } + if (read_bytes_ > lease_.io_buffer().size_) { + ready_.emplace(absl::OutOfRangeError( + "record exceeds registered read buffer capacity")); + return; } - - // Only the disk read suspends while holding the BlockState pointer, so it - // is the only path that has to keep the state alive with a pin. ++state->pins_; - struct PinGuard { - WorkerStore* store_ = nullptr; - BlockState* state_ = nullptr; - ~PinGuard() { - if (state_ == nullptr) { - return; - } - --state_->pins_; - if (state_->pins_ == 0 && state_->release_pending_) { - ReleaseStagingBuffer(*store_, *state_); - } - } - } pin{&store, state}; + pinned_ = state; + if (trace_) trace_->disk_read_ = true; +} - if (trace != nullptr) { - trace->disk_read_ = true; - } - FixedBuffer io = lease.io_buffer(); - if (read_bytes > io.size_) { - co_return absl::Status(absl::StatusCode::kOutOfRange, - "record exceeds registered read buffer capacity"); - } - FixedBuffer record_buffer = io; - record_buffer.size_ = read_bytes; - if (trace != nullptr) { - trace->io_submit_ns_ = ReadTraceNowNanos(); - } - auto read = co_await ReadStorageBuffer(*store.worker_, store.files_[file_id], - record_buffer, lease.registered(), - aligned_offset); - if (trace != nullptr) { - trace->io_complete_ns_ = ReadTraceNowNanos(); - } - if (!read.ok()) { - co_return read.status(); - } - if (*read != read_bytes) { - co_return absl::Status(absl::StatusCode::kInternal, - "short compact record read"); +bool StorageEngine::Impl::LoadValueAwaitable::PopulationChanged() const { + return partition_ && (partition_->replication_epoch_ != replication_epoch_ || + engine_.EffectiveRecordDbEpoch(*partition_, db_id_) != + *expected_db_epoch_); +} + +void StorageEngine::Impl::LoadValueAwaitable::Retry(absl::Status failure) { + assert(partition_ && absl::IsAborted(failure)); + Unpin(); + lease_.Reset(); + ready_.reset(); + // Inline reads have no extent manifest. Seed the failure so the slow path + // resolves the index before any further I/O and preserves corruption errors + // when the index still endorses the failed physical record. + slow_ = engine_.LoadValueSlow(store_, *partition_, db_id_, key_, digest_, + location_, {}, trace_, std::move(failure)); +} + +void StorageEngine::Impl::LoadValueAwaitable::FinishBufferWait() { + auto acquired = buffer_wait_->await_resume(); + buffer_wait_.reset(); + if (PopulationChanged()) { + ready_.emplace(absl::NotFoundError("key not found")); + } else if (!acquired.ok()) { + ready_.emplace(acquired.status()); + } else { + lease_ = std::move(*acquired); + PrepareLocal(); + if (partition_ && ready_ && absl::IsAborted(ready_->status())) + Retry(ready_->status()); } +} + +void StorageEngine::Impl::LoadValueAwaitable::BufferReady( + bycorf::RemoteWork* work) { + auto& read = *static_cast(work); + read.FinishBufferWait(); + auto next = read.SuspendRead(read.awaiting_); + if (next != std::noop_coroutine()) read.store_.worker_->Enqueue(next); +} +std::coroutine_handle<> StorageEngine::Impl::LoadValueAwaitable::await_suspend( + std::coroutine_handle<> awaiting) { + if (buffer_wait_) { + awaiting_ = awaiting; + run_fn_ = &BufferReady; + task_class_ = store_.worker_->IsBackground(awaiting) + ? bycorf::TaskClass::kBackground + : bycorf::CurrentTaskClass(); + if (buffer_wait_->SuspendWith(*this)) return std::noop_coroutine(); + FinishBufferWait(); + } + return SuspendRead(awaiting); +} + +std::coroutine_handle<> StorageEngine::Impl::LoadValueAwaitable::SuspendRead( + std::coroutine_handle<> awaiting) { + if (slow_.valid()) + return std::move(slow_).operator co_await().await_suspend(awaiting); + if (ready_) return awaiting; + auto buffer = lease_.io_buffer(); + buffer.size_ = read_bytes_; + if (trace_) trace_->io_submit_ns_ = ReadTraceNowNanos(); + auto status = + lease_.registered() + ? store_.worker_->SubmitReadFixed(store_.files_[file_id_], buffer, + aligned_offset_, this) + : store_.worker_->SubmitRead(store_.files_[file_id_], + std::span(buffer.data_, buffer.size_), + aligned_offset_, this); + if (Suspend(awaiting, status)) return std::noop_coroutine(); + if (trace_) trace_->io_complete_ns_ = ReadTraceNowNanos(); + if (partition_ && absl::IsAborted(status)) { + Retry(std::move(status)); + return std::move(slow_).operator co_await().await_suspend(awaiting); + } + return awaiting; +} + +void StorageEngine::Impl::LoadValueAwaitable::Complete(bycorf::Worker& worker, + int bytes, + unsigned flags) { + if (trace_) trace_->io_complete_ns_ = ReadTraceNowNanos(); + if (bytes >= 0) worker.RecordStorageReadCompletion(bytes); + if (partition_ && bytes >= 0 && + static_cast(bytes) == read_bytes_) { + // Only metadata is examined on the completion poller. CRC and large-value + // copies stay in the caller's scheduled slice. A stale identity must select + // the retry continuation before the caller's co_await can finish. + auto parsed = ParseRecord(lease_.io_buffer().data_ + record_headroom_); + if (parsed.ok()) + payload_ = *parsed; + else if (absl::IsAborted(parsed.status()) && !PopulationChanged()) { + Retry(parsed.status()); + auto child = + std::move(slow_).operator co_await().await_suspend(awaiting_); + // CQEs are polled outside the parent's task class. Preserve background + // scheduling when a backup/maintenance read starts its retry here. + if (worker.IsBackground(awaiting_)) worker.RegisterBackground(child); + worker.Enqueue(child); + return; + } else + ready_.emplace(parsed.status()); + } + OneShotIoAwaitable::Complete(worker, bytes, flags); +} + +StorageEngine::Impl::LoadValueAwaitable::Result +StorageEngine::Impl::LoadValueAwaitable::await_resume() { + // Fence successful reads here, after the scheduling gap following the CQE. + // Completion only needs this check when deciding whether to start a retry. + if (PopulationChanged()) return absl::NotFoundError("key not found"); + if (slow_.valid()) return std::move(slow_).TakeResult(); + if (ready_) return std::move(*ready_); + auto read = + Resume(lease_.registered() ? "fixed-buffer read failed" : "read failed"); + if (!read.ok()) return read.status(); + if (static_cast(*read) != read_bytes_) + return absl::InternalError("short compact record read"); + const auto* data = lease_.io_buffer().data_ + record_headroom_; + if (!payload_) { + auto parsed = ParseRecord(data); + if (!parsed.ok()) return parsed.status(); + payload_ = *parsed; + } + return FinishRecord(data, *payload_, false); +} + +absl::StatusOr +StorageEngine::Impl::LoadValueAwaitable::ParseRecord( + const std::byte* data) const { RecordHeader record{}; std::string_view disk_key; - const std::byte* record_data = io.data_ + record_headroom; - std::span record_bytes(record_data, - location.total_disk_bytes()); - if (!DecodeRecordHeader(record_bytes, &record, &disk_key) || - record.db_id_ != db_id || (!record.key_indirect_ && disk_key != key) || + if (!DecodeRecordHeader(std::span(data, location_.total_disk_bytes()), + &record, &disk_key) || + record.db_id_ != db_id_ || (!record.key_indirect_ && disk_key != key_) || record.kind_ != RecordKind::kValue || - record.db_epoch_ != expected_db_epoch.value_or(DbEpoch(db_id)) || - record.mutation_sequence_ != location.mutation_sequence_ || - record.replication_epoch_ != replication_epoch || - record.allocation_epoch_ != location.allocation_epoch() || - record.expire_at_ms_ != location.expire_at_ms_ || - record.value_type_ != location.value_type() || - record.external_ != location.external() || - record.key_indirect_ != location.key_indirect() || - record.logical_size_ != location.logical_size_ || - record.total_disk_bytes_ != location.total_disk_bytes()) { - co_return absl::Status(absl::StatusCode::kAborted, - "record does not match in-memory location"); - } - const std::byte* payload_data = record_data + record.header_bytes_; - if (record.key_indirect_ && record.key_bytes_ != key.size()) { - co_return absl::Status(absl::StatusCode::kAborted, - "record key does not match location"); - } - const std::size_t value_bytes = record.payload_bytes_; + record.db_epoch_ != (expected_db_epoch_ ? *expected_db_epoch_ + : engine_.DbEpoch(db_id_)) || + record.mutation_sequence_ != location_.mutation_sequence_ || + record.replication_epoch_ != replication_epoch_ || + record.allocation_epoch_ != location_.allocation_epoch() || + record.expire_at_ms_ != location_.expire_at_ms_ || + record.value_type_ != location_.value_type() || + record.external_ != location_.external() || + record.key_indirect_ != location_.key_indirect() || + record.logical_size_ != location_.logical_size_ || + record.total_disk_bytes_ != location_.total_disk_bytes()) { + return absl::AbortedError("record does not match in-memory location"); + } + if (record.key_indirect_ && record.key_bytes_ != key_.size()) + return absl::AbortedError("record key does not match location"); if (record.value_type_ == ValueType::kString && !record.grouped_ && - !record.auxiliary_group_ && value_bytes != record.logical_size_) { - co_return absl::Status(absl::StatusCode::kInternal, - "inline value length does not match metadata"); - } - if (Crc32c(std::span(payload_data, record.payload_bytes_)) != - record.payload_checksum_) { - co_return absl::Status(absl::StatusCode::kInternal, - "record value checksum mismatch"); - } - const std::byte* framed_value = payload_data; - if (value_bytes > DirectGetValueLimit()) { - std::memmove(io.data_, payload_data, value_bytes); - framed_value = io.data_; - } - if (trace != nullptr) { - trace->decode_done_ns_ = ReadTraceNowNanos(); - } - const std::size_t value_offset = - static_cast(framed_value - lease.bytes().data()); - co_return LoadedValue{std::move(lease), value_offset, value_bytes, - record.txid_}; + !record.auxiliary_group_ && record.payload_bytes_ != record.logical_size_) + return absl::InternalError("inline value length does not match metadata"); + return Payload{record.txid_, record.header_bytes_, record.payload_bytes_, + record.payload_checksum_}; +} + +StorageEngine::Impl::LoadValueAwaitable::Result +StorageEngine::Impl::LoadValueAwaitable::FinishRecord(const std::byte* data, + Payload payload, + bool staged) { + const auto* value = data + payload.offset_; + if (Crc32c(std::span(value, payload.bytes_)) != payload.checksum_) + return absl::InternalError("record value checksum mismatch"); + if (staged || payload.bytes_ > engine_.DirectGetValueLimit()) { + std::memmove(lease_.io_buffer().data_, value, payload.bytes_); + value = lease_.io_buffer().data_; + } + const auto offset = static_cast(value - lease_.bytes().data()); + if (trace_) trace_->decode_done_ns_ = ReadTraceNowNanos(); + return LoadedValue{std::move(lease_), offset, payload.bytes_, payload.txid_}; } } // namespace lavik::storage diff --git a/tests/grouped/recovery_e2e_test.cpp b/tests/grouped/recovery_e2e_test.cpp index a0458dee0..0622ff957 100644 --- a/tests/grouped/recovery_e2e_test.cpp +++ b/tests/grouped/recovery_e2e_test.cpp @@ -1170,7 +1170,10 @@ TEST(GroupedRecoveryE2e, CheckpointAccountsForReassignedPhysicalOwners) { EXPECT_NE(recovered.Log().find("loaded shutdown checkpoint generation="), std::string::npos) << recovered.Log(); - EXPECT_EQ(recovered.Log().find("falling back"), std::string::npos) + // A kernel may independently fall back from provided-buffer receive. + // Only checkpoint recovery falling back to a record scan is a failure. + EXPECT_EQ(recovered.Log().find("falling back to record scan"), + std::string::npos) << recovered.Log(); EXPECT_EQ(recovered.Command({"HGET", "hash", "field"}), value); EXPECT_EQ(recovered.Command({"SET", "string", value}), "+OK"); diff --git a/tests/load_value_awaitable_test.cpp b/tests/load_value_awaitable_test.cpp new file mode 100644 index 000000000..d5e1a4302 --- /dev/null +++ b/tests/load_value_awaitable_test.cpp @@ -0,0 +1,487 @@ +/* + * Copyright (C) 2026 EloqData Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include "../src/storage/engine/impl.h" +#include "gtest/gtest.h" + +namespace lavik::storage { + +// Deliver deterministic CQEs for identity races and real I/O for buffer +// handoff. +class LoadValueAwaitableTestPeer : public testing::Test { + protected: + using Awaitable = StorageEngine::Impl::LoadValueAwaitable; + using Result = Awaitable::Result; + using ExtentRead = StorageEngine::Impl::ReadExtentAwaitable; + + void SetUp() override { + saved_context_ = std::move(bycorf::MutableThisWorker()); + wake_fd_ = eventfd(0, EFD_CLOEXEC | EFD_NONBLOCK); + ASSERT_GE(wake_fd_, 0); + cross_core_.mailbox(0).wake_fd_ = wake_fd_; + worker_.BindCrossCore(0, &cross_core_); + bycorf::SetThisWorker(0, &cross_core_, &worker_); + ASSERT_TRUE(worker_.Init({.recv_buffer_count_ = 0}).ok()); + auto& engine = *storage_.impl_; + engine.worker_count_ = 1; + engine.stores_.push_back( + std::make_unique()); + auto& store = *engine.stores_[0]; + store.worker_ = &worker_; + store.partitions_.resize(1); + store.partition_indexes_.resize(kLogicalDatabaseCount); + store.partitions_[0].indexes_ = store.partition_indexes_; + store.partitions_[0].replication_epoch_ = 1; + store.files_.push_back(bycorf::FixedFile{}); + engine.db_epochs_[0].store(1); + engine.devices_.push_back(StorageDevice{.capacity_blocks_ = 2}); + engine.block_state_lookup_.push_back({&block_, 1, 2}); + block_.Reset(0, 7); + block_.allocated_ = true; + + RegisteredBufferPoolOptions options; + options.storage_write_buffer_count_ = 1; + options.registered_bytes_ = + options.write_buffer_bytes_ + options.read_slot_bytes_; + ASSERT_TRUE(Pool().Init(worker_, options).ok()); + ASSERT_EQ(Pool().available_read_buffers(), 1); + + header_.header_bytes_ = RecordHeaderBytes(key_.size()); + header_.key_bytes_ = key_.size(); + header_.value_type_ = ValueType::kString; + header_.logical_size_ = value_.size(); + header_.payload_bytes_ = value_.size(); + header_.total_disk_bytes_ = + AlignRecord(header_.header_bytes_ + value_.size()); + header_.allocation_epoch_ = 7; + header_.mutation_sequence_ = 9; + header_.payload_checksum_ = Crc32c(std::as_bytes(std::span(value_))); + const auto metadata = RecordLocation::PackedMetadata::Encode( + kBlockHeaderBytes + 8, header_.total_disk_bytes_, 0, false, false, + false, false, false, false, RecordKind::kValue, ValueType::kString); + location_ = RecordLocation(1, 9, 7, 0, value_.size(), metadata); + } + + void TearDown() override { + worker_.Shutdown(); + storage_.impl_.reset(); + if (wake_fd_ >= 0) close(wake_fd_); + bycorf::MutableThisWorker() = std::move(saved_context_); + } + + auto& Store() { return *storage_.impl_->stores_[0]; } + RegisteredBufferPool& Pool() { return Store().buffers_; } + auto Read() { + return storage_.impl_->LoadValue(Store(), Store().partitions_[0], 0, key_, + digest_, location_, {}); + } + auto ReadLocal() { + return storage_.impl_->LoadValueLocal(Store(), 0, key_, location_, 1); + } + bool HasHelper(const Awaitable& read) { return read.slow_.valid(); } + + void Encode(std::byte* data) { + ASSERT_TRUE(EncodeRecordHeader(header_, key_, + std::span(data, header_.header_bytes_))); + std::memcpy(data + header_.header_bytes_, value_.data(), value_.size()); + } + void Complete(Awaitable& read, int bytes = -1) { + ASSERT_NE(read.pinned_, nullptr); + Encode(read.lease_.io_buffer().data_ + read.record_headroom_); + read.Complete(worker_, bytes == -1 ? read.read_bytes_ : bytes, 0); + } + void Stage() { + location_.set_in_memory(true); + block_.in_memory_ = true; + if (block_.staging_slot_ == 0) { + block_.staging_slot_ = Store().staging_slots_.size(); + Store().staging_slots_.emplace_back(); + ASSERT_TRUE(Pool().TryAcquireWriteBuffer( + &Store().staging_slots_.back().write_buffer_id_)); + } + auto buffer = storage_.impl_->StagingBufferFor(Store(), block_); + Encode(buffer.data_ + location_.record_offset()); + } + void IndexLocation() { + ASSERT_NE(Store() + .partitions_[0] + .indexes_[0] + .InsertOrAssign(digest_, key_, location_) + .entry_, + nullptr); + } + void Invalidate(bool database) { + if (database) + storage_.impl_->db_epochs_[0].fetch_add(1); + else + ++Store().partitions_[0].replication_epoch_; + } + void CandidatePopulation() { + storage_.impl_->replica_loading_.store(true); + auto& partition = Store().partitions_[0]; + partition.replica_sync_ = std::make_unique< + StorageEngine::Impl::WorkerStore::PartitionStore::ReplicaSyncState>(); + partition.replica_sync_->local_db_epochs_[0] = 4; + header_.db_epoch_ = 4; + } + void ExpectValue(const Result& result) { + ASSERT_TRUE(result.ok()) << result.status(); + auto bytes = result->value(); + EXPECT_EQ(std::string_view(reinterpret_cast(bytes.data()), + bytes.size()), + value_); + } + + // Simulate only the submission, retaining the real completion -> retry -> + // parent continuation chain, with an owned Task to catch early/double resume. + Task AwaitCompletion(Awaitable& read) { + struct Arm { + Awaitable& read_; + bool await_ready() const { return false; } + void await_suspend(std::coroutine_handle<> parent) { + read_.awaiting_ = parent; + } + void await_resume() const {} + }; + co_await Arm{read}; + co_return read.await_resume(); + } + Task AwaitExisting(Awaitable& read) { co_return co_await read; } + Task AwaitExtent(ExtentRead& read) { co_return co_await read; } + bool BufferWaitIsBackground(Awaitable& read) { + return read.task_class_ == bycorf::TaskClass::kBackground; + } + auto ExtentSlice(ExtentRef ref, std::size_t offset, + std::span destination) { + return storage_.impl_->ReadExtentSlice(Store(), ref, 0, offset, + destination); + } + + // Use a real fixed-file read for pool -> I/O handoff, rather than completing + // the request before the continuation actually submits it. + void InstallDiskImage(std::span bytes) { + char path[] = "/tmp/lavik-awaitable-XXXXXX"; + const int fd = mkstemp(path); + ASSERT_GE(fd, 0); + const auto offset = storage_.impl_->FileOffset(1).second; + const auto written = pwrite(fd, bytes.data(), bytes.size(), offset); + close(fd); + ASSERT_EQ(written, static_cast(bytes.size())); + struct Opened : bycorf::IoCompletion { + bool done_ = false; + int result_ = -1; + void Complete(bycorf::Worker&, int result, unsigned) override { + result_ = result; + done_ = true; + } + } opened; + ASSERT_TRUE(worker_.RegisterFixedFiles(1).ok()); + auto submitted = + worker_.SubmitOpenDirect(path, O_RDONLY | O_DIRECT, 0, {0}, &opened); + if (submitted.ok()) { + const auto until = + std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!opened.done_ && std::chrono::steady_clock::now() < until) + worker_.RunOnce(false); + } + unlink(path); + ASSERT_TRUE(submitted.ok()) << submitted; + ASSERT_TRUE(opened.done_); + ASSERT_EQ(opened.result_, 0); + } + void InstallRecord() { + std::vector bytes( + AlignDirect(location_.record_offset() + location_.total_disk_bytes())); + Encode(bytes.data() + location_.record_offset()); + InstallDiskImage(bytes); + } + ExtentRef ExtentReference() { + return {1, 7, static_cast(value_.size()), + Crc32c(std::as_bytes(std::span(value_)))}; + } + void InstallExtent(bool corrupt = false) { + block_.kind_ = BlockKind::kPayloadExtent; + auto ref = ExtentReference(); + std::vector bytes( + AlignDirect(kBlockHeaderBytes + value_.size())); + BlockHeader header; + header.block_id_ = 1; + header.allocation_epoch_ = 7; + header.kind_ = BlockKind::kPayloadExtent; + header.layout_worker_count_ = 1; + header.committed_bytes_ = kBlockHeaderBytes + value_.size(); + header.extent_payload_bytes_ = value_.size(); + header.extent_payload_checksum_ = ref.payload_checksum_; + EncodeBlockHeader(header, std::span( + bytes.data(), kBlockHeaderSlotBytes)); + std::memcpy(bytes.data() + kBlockHeaderBytes, value_.data(), value_.size()); + // Corrupt beyond the slice returned by the test, inside the actual payload. + if (corrupt) bytes[kBlockHeaderBytes + value_.size() - 1] ^= std::byte{1}; + InstallDiskImage(bytes); + } + + template + auto Start(Task& task, bool background = false) { + auto handle = std::move(task).operator co_await().await_suspend( + std::noop_coroutine()); + if (background) worker_.RegisterBackground(handle); + handle.resume(); + return handle; + } + template + void Drain(Task& task) { + const auto until = + std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!task.done() && std::chrono::steady_clock::now() < until) + worker_.RunOnce(false); + ASSERT_TRUE(task.done()); + } + bool HelperIsBackground(Awaitable& read) { + return worker_.IsBackground( + std::move(read.slow_).operator co_await().handle_); + } + + bycorf::CurrentWorker saved_context_; + bycorf::CrossCore cross_core_{1}; + int wake_fd_ = -1; + bycorf::Worker worker_; + BlockState block_; + StorageEngine storage_{StorageEngineOptions{}}; + std::string key_ = "direct-key"; + std::string value_ = "direct value"; + Digest digest_{17}; + RecordHeader header_; + RecordLocation location_; +}; + +TEST_F(LoadValueAwaitableTestPeer, BufferWaitersPreserveFifoHandoff) { + ASSERT_NO_FATAL_FAILURE(InstallRecord()); + auto held = Pool().AcquireReadBuffer().await_resume(); + ASSERT_TRUE(held.ok()); + auto first = Read(); + auto second = Read(); + auto first_task = AwaitExisting(first); + auto second_task = AwaitExisting(second); + Start(first_task, true); + Start(second_task); + EXPECT_TRUE(BufferWaitIsBackground(first)); + EXPECT_EQ(block_.pins_, 0); + held->Reset(); + // Handoff schedules I/O; it must not resume inside the release stack. + EXPECT_FALSE(first_task.done()); + EXPECT_EQ(Pool().available_read_buffers(), 0); + ASSERT_NO_FATAL_FAILURE(Drain(first_task)); + EXPECT_FALSE(second_task.done()); + // Taking/destroying the first result hands its sole slot to the second. + ExpectValue(std::move(first_task).TakeResult()); + ASSERT_NO_FATAL_FAILURE(Drain(second_task)); + ExpectValue(std::move(second_task).TakeResult()); + EXPECT_FALSE(HasHelper(first)); + EXPECT_FALSE(HasHelper(second)); + EXPECT_EQ(Pool().available_read_buffers(), 1); +} + +TEST_F(LoadValueAwaitableTestPeer, PopulationChangeDuringBufferWaitSkipsIo) { + auto held = Pool().AcquireReadBuffer().await_resume(); + ASSERT_TRUE(held.ok()); + auto read = Read(); + auto task = AwaitExisting(read); + Start(task); + Invalidate(true); + held->Reset(); + ASSERT_NO_FATAL_FAILURE(Drain(task)); + EXPECT_TRUE(absl::IsNotFound(std::move(task).TakeResult().status())); + EXPECT_FALSE(HasHelper(read)); + EXPECT_EQ(block_.pins_, 0); + EXPECT_EQ(Pool().available_read_buffers(), 1); +} + +TEST_F(LoadValueAwaitableTestPeer, RelocationDuringBufferWaitUsesNewLocation) { + auto held = Pool().AcquireReadBuffer().await_resume(); + ASSERT_TRUE(held.ok()); + auto read = Read(); + auto task = AwaitExisting(read); + Start(task, true); + EXPECT_EQ(block_.pins_, 0); + // No physical pointer was retained while waiting: recycle before handoff. + block_.Reset(0, 8); + block_.allocated_ = true; + header_.allocation_epoch_ = 8; + location_ = RecordLocation(1, 9, 8, 0, value_.size(), location_.metadata_); + Stage(); + IndexLocation(); + held->Reset(); + ASSERT_NO_FATAL_FAILURE(Drain(task)); + ExpectValue(std::move(task).TakeResult()); + EXPECT_TRUE(HasHelper(read)); + EXPECT_EQ(block_.pins_, 0); + EXPECT_EQ(Pool().available_read_buffers(), 1); +} + +TEST_F(LoadValueAwaitableTestPeer, ExtentWaitRetainsPinAndReturnsSlice) { + ASSERT_NO_FATAL_FAILURE(InstallExtent()); + auto held = Pool().AcquireReadBuffer().await_resume(); + ASSERT_TRUE(held.ok()); + std::array output{}; + { + auto read = ExtentSlice(ExtentReference(), 2, output); + EXPECT_FALSE(read.await_ready()); + EXPECT_EQ(block_.pins_, 1); + auto task = AwaitExtent(read); + Start(task, true); + held->Reset(); + EXPECT_FALSE(task.done()); + ASSERT_NO_FATAL_FAILURE(Drain(task)); + auto status = std::move(task).TakeResult(); + EXPECT_TRUE(status.ok()) << status; + EXPECT_EQ(std::string_view(reinterpret_cast(output.data()), + output.size()), + std::string_view(value_).substr(2, output.size())); + } + EXPECT_EQ(block_.pins_, 0); + EXPECT_EQ(Pool().available_read_buffers(), 1); +} + +TEST_F(LoadValueAwaitableTestPeer, ExtentSliceChecksWholePayloadBeforeCopying) { + ASSERT_NO_FATAL_FAILURE(InstallExtent(true)); + std::array output{std::byte{0x55}, std::byte{0x55}}; + { + auto read = ExtentSlice(ExtentReference(), 0, output); + auto task = AwaitExtent(read); + Start(task); + ASSERT_NO_FATAL_FAILURE(Drain(task)); + auto status = std::move(task).TakeResult(); + EXPECT_TRUE(absl::IsInternal(status)) << status; + EXPECT_EQ(status.message(), "extent payload checksum mismatch"); + EXPECT_EQ(output[0], std::byte{0x55}); + EXPECT_EQ(output[1], std::byte{0x55}); + } + EXPECT_EQ(block_.pins_, 0); + EXPECT_EQ(Pool().available_read_buffers(), 1); +} + +TEST_F(LoadValueAwaitableTestPeer, CompletionTransfersLeaseAndReleasesPin) { + std::optional result; + { + auto read = Read(); + ASSERT_FALSE(read.await_ready()); + ASSERT_NO_FATAL_FAILURE(Complete(read)); + result = read.await_resume(); + ExpectValue(*result); + } + EXPECT_EQ(block_.pins_, 0); + EXPECT_EQ(Pool().available_read_buffers(), 0); + result.reset(); + EXPECT_EQ(Pool().available_read_buffers(), 1); +} + +TEST_F(LoadValueAwaitableTestPeer, PopulationChangesDiscardCompletedRead) { + for (bool database : {false, true}) { + // Exercise valid/mismatched metadata at completion and the scheduling gap. + for (int timing : {0, 1, 2}) { + SCOPED_TRACE(testing::Message() + << "database=" << database << " timing=" << timing); + { + auto read = Read(); + if (timing != 2) Invalidate(database); + if (timing == 1) ++header_.allocation_epoch_; + ASSERT_NO_FATAL_FAILURE(Complete(read)); + if (timing == 1) --header_.allocation_epoch_; + if (timing == 2) Invalidate(database); + EXPECT_FALSE(HasHelper(read)); + EXPECT_TRUE(absl::IsNotFound(read.await_resume().status())); + } + EXPECT_EQ(block_.pins_, 0); + EXPECT_EQ(Pool().available_read_buffers(), 1); + if (database) + ++header_.db_epoch_; + else + ++header_.replication_epoch_; + } + } +} + +TEST_F(LoadValueAwaitableTestPeer, CandidatePopulationUsesMappedEpoch) { + CandidatePopulation(); + auto read = Read(); + EXPECT_FALSE(HasHelper(read)); + ASSERT_NO_FATAL_FAILURE(Complete(read)); + ExpectValue(read.await_resume()); +} + +TEST_F(LoadValueAwaitableTestPeer, ErrorsReleaseResources) { + for (int completion : {0, -EIO, 4096}) { + { + auto read = Read(); + ASSERT_FALSE(read.await_ready()); + header_.payload_checksum_ ^= 1; + ASSERT_NO_FATAL_FAILURE(Complete(read, completion)); + header_.payload_checksum_ ^= 1; + EXPECT_FALSE(read.await_resume().ok()); + } + EXPECT_EQ(block_.pins_, 0); + EXPECT_EQ(Pool().available_read_buffers(), 1); + } +} + +TEST_F(LoadValueAwaitableTestPeer, LocalReadReturnsStaleLocationToItsCaller) { + block_.freeing_ = true; + auto read = ReadLocal(); + ASSERT_TRUE(read.await_ready()); + EXPECT_FALSE(HasHelper(read)); + EXPECT_TRUE(absl::IsAborted(read.await_resume().status())); +} + +TEST_F(LoadValueAwaitableTestPeer, + RelocationRetriesBeforeResumingBackgroundCaller) { + auto read = Read(); + auto task = AwaitCompletion(read); + Start(task, true); + ++header_.allocation_epoch_; + ASSERT_NO_FATAL_FAILURE(Complete(read)); + EXPECT_FALSE(task.done()); + EXPECT_TRUE(HelperIsBackground(read)); + EXPECT_EQ(block_.pins_, 0); + EXPECT_EQ(Pool().available_read_buffers(), 1); + --header_.allocation_epoch_; + // The new physical record has the same logical mutation and lives in staging. + location_.metadata_ = RecordLocation::PackedMetadata::Encode( + location_.record_offset() + 512, location_.total_disk_bytes(), 0, false, + false, false, false, false, false, RecordKind::kValue, + ValueType::kString); + Stage(); + IndexLocation(); + ASSERT_NO_FATAL_FAILURE(Drain(task)); + ExpectValue(std::move(task).TakeResult()); +} + +TEST_F(LoadValueAwaitableTestPeer, UnchangedBadLocationIsCorruption) { + IndexLocation(); + auto read = Read(); + auto task = AwaitCompletion(read); + Start(task); + ++header_.allocation_epoch_; + ASSERT_NO_FATAL_FAILURE(Complete(read)); + ASSERT_NO_FATAL_FAILURE(Drain(task)); + EXPECT_TRUE(absl::IsInternal(std::move(task).TakeResult().status())); +} + +TEST_F(LoadValueAwaitableTestPeer, ConcurrentLogicalMutationIsNotRelocation) { + auto read = Read(); + auto task = AwaitCompletion(read); + Start(task); + ++header_.mutation_sequence_; + ASSERT_NO_FATAL_FAILURE(Complete(read)); + ++location_.mutation_sequence_; + IndexLocation(); + ASSERT_NO_FATAL_FAILURE(Drain(task)); + EXPECT_TRUE(absl::IsNotFound(std::move(task).TakeResult().status())); +} + +} // namespace lavik::storage