From 02eb51432e4273e5e4ccbb592a287b6480e1e90b Mon Sep 17 00:00:00 2001 From: Aoi Date: Mon, 19 Oct 2026 18:26:47 +0800 Subject: [PATCH] [Store] Remove snapshot serializers orphaned by the legacy snapshot removal #4406 removed the Master-generated snapshot and restore path, but left behind the code that only that path called. None of it has a caller outside its own tests: - serialize/serializer.{h,cpp}: the msgpack Serializer specializations for allocators, buffers, replicas and mounted segments, SerializationHelper and SerializationError. - ha/snapshot: StoreResourceSnapshotCodec, AllocatorSnapshotCodec and LocalSsdCodec. - SegmentSerializer and SegmentView in segment.h, TaskManagerSerializer with the restore_task/clear_all helpers only it used, and LocalSsdManager::ExportPersistedState/RestorePersistedState with LocalSsdPersistedState. - The detached snapshot path added by #3706: SegmentPool::CaptureSnapshot and RestoreSnapshot, OffsetBufferAllocator and OffsetAllocator CaptureSnapshot/Restore with their snapshot structs and layout validation, AdoptMode::ReplacePool and IsKnownSegmentStatus. SegmentPool::WriteAccess::PrepareAdopt stays in its pre-#3706 form: the read-lock test still uses it to install a blocking allocator. The OffsetBufferAllocator node cap goes back to the literal 64Mi it had before #3706 named it OffsetAllocatorSnapshot::kMaxNodes. DestroyedPoolAndRetainedAllocatorRemoveMetricLabels stays, rewritten to mount directly instead of restoring a snapshot. #3706 also made buffer allocator destructors remove their segment's metric labels and made remove_segment_metrics keep a label that still has usage; both are on the ordinary mount path and lost their only coverage with the snapshot tests. --- mooncake-store/include/allocation_strategy.h | 1 - mooncake-store/include/allocator.h | 37 - .../ha/snapshot/allocator_snapshot_codec.h | 22 - .../include/ha/snapshot/local_ssd_codec.h | 20 - .../snapshot/store_resource_snapshot_codec.h | 39 - mooncake-store/include/local_ssd/manager.h | 3 - .../include/local_ssd/persisted_state.h | 22 - .../offset_allocator/offset_allocator.h | 73 -- mooncake-store/include/replica.h | 3 +- mooncake-store/include/segment.h | 48 +- mooncake-store/include/segment/pool.h | 13 - .../include/segment/pool_write_access.h | 6 +- mooncake-store/include/segment/snapshot.h | 26 - mooncake-store/include/segment/status.h | 14 - .../include/segment_allocator_registration.h | 5 - mooncake-store/include/serialize/serializer.h | 150 --- mooncake-store/include/task_manager.h | 22 - mooncake-store/include/types.h | 11 - mooncake-store/src/CMakeLists.txt | 5 - mooncake-store/src/allocator.cpp | 77 +- .../ha/snapshot/allocator_snapshot_codec.cpp | 97 -- .../src/ha/snapshot/local_ssd_codec.cpp | 153 --- .../store_resource_snapshot_codec.cpp | 388 ------ mooncake-store/src/local_ssd/manager.cpp | 40 - mooncake-store/src/master_service.cpp | 1 - mooncake-store/src/offset_allocator.cpp | 311 ----- mooncake-store/src/segment.cpp | 454 ------- mooncake-store/src/segment/pool_mount.cpp | 16 +- mooncake-store/src/segment/snapshot.cpp | 123 -- mooncake-store/src/serialize/serializer.cpp | 1080 ----------------- mooncake-store/src/task_manager.cpp | 189 --- mooncake-store/tests/CMakeLists.txt | 4 - .../tests/buffer_allocator_test.cpp | 118 -- .../ha/snapshot/local_ssd_codec_test.cpp | 79 -- .../store_resource_snapshot_codec_test.cpp | 881 -------------- .../tests/local_ssd/local_ssd_test.cpp | 26 - mooncake-store/tests/master_metrics_test.cpp | 76 -- .../tests/offset_allocator_test.cpp | 334 ----- mooncake-store/tests/segment_pool_test.cpp | 495 +------- mooncake-store/tests/segment_test.cpp | 52 - mooncake-store/tests/serializer_test.cpp | 225 ---- mooncake-store/tests/task_manager_test.cpp | 148 --- 42 files changed, 24 insertions(+), 5863 deletions(-) delete mode 100644 mooncake-store/include/ha/snapshot/allocator_snapshot_codec.h delete mode 100644 mooncake-store/include/ha/snapshot/local_ssd_codec.h delete mode 100644 mooncake-store/include/ha/snapshot/store_resource_snapshot_codec.h delete mode 100644 mooncake-store/include/local_ssd/persisted_state.h delete mode 100644 mooncake-store/include/segment/snapshot.h delete mode 100644 mooncake-store/include/serialize/serializer.h delete mode 100644 mooncake-store/src/ha/snapshot/allocator_snapshot_codec.cpp delete mode 100644 mooncake-store/src/ha/snapshot/local_ssd_codec.cpp delete mode 100644 mooncake-store/src/ha/snapshot/store_resource_snapshot_codec.cpp delete mode 100644 mooncake-store/src/segment/snapshot.cpp delete mode 100644 mooncake-store/src/serialize/serializer.cpp delete mode 100644 mooncake-store/tests/ha/snapshot/local_ssd_codec_test.cpp delete mode 100644 mooncake-store/tests/ha/snapshot/store_resource_snapshot_codec_test.cpp delete mode 100644 mooncake-store/tests/serializer_test.cpp diff --git a/mooncake-store/include/allocation_strategy.h b/mooncake-store/include/allocation_strategy.h index 762dff87a7..7888649fac 100644 --- a/mooncake-store/include/allocation_strategy.h +++ b/mooncake-store/include/allocation_strategy.h @@ -204,7 +204,6 @@ class AllocatorManager { allocators_; std::unordered_map owner_by_name_; friend class ScopedSegmentAccess; - friend class SegmentSerializer; }; /** diff --git a/mooncake-store/include/allocator.h b/mooncake-store/include/allocator.h index 7a45796faa..58f3441a74 100644 --- a/mooncake-store/include/allocator.h +++ b/mooncake-store/include/allocator.h @@ -171,7 +171,6 @@ class AllocatedBuffer { std::optional offset_handle_{ std::nullopt}; - friend class Serializer; friend class Replica; friend class SegmentAllocatorRegistration; }; @@ -349,26 +348,6 @@ std::optional ImportCachelibBufferAllocator( const std::vector& allocations, ReplicaType replica_type = ReplicaType::MEMORY); -// A detached capture of host-memory allocator metadata; owns no live allocator -// and contributes neither usage nor capacity metrics. -struct OffsetBufferAllocatorSnapshot { - std::string segment_name; - size_t base; - size_t capacity; - size_t used_bytes; - std::string transport_endpoint; - offset_allocator::OffsetAllocatorSnapshot allocation_state; - - // Validate the detached snapshot without publishing resources or metrics. - [[nodiscard]] tl::expected Validate() const; - - private: - tl::expected ValidateMetadata() const; - // Restore composes this cheap outer check with OffsetAllocator::Restore's - // mandatory layout validation, avoiding two graph scans in one factory. - friend class OffsetBufferAllocator; -}; - /** * OffsetBufferAllocator manages memory allocation using the OffsetAllocator * strategy, which provides efficient memory allocation with bin-based @@ -384,18 +363,6 @@ class OffsetBufferAllocator ~OffsetBufferAllocator() override; - // Restore an unpublished allocator from decoded state. Installs the - // allocation layout and accounts usage before returning a runtime-ready - // object; capacity accounting remains the owning pool's responsibility. - // Consumes the snapshot and rejects inconsistent state with - // ErrorCode::INVALID_PARAMS. - static tl::expected, ErrorCode> - Restore(OffsetBufferAllocatorSnapshot snapshot); - - // Requires a forked snapshot child or externally quiesced state. Does not - // acquire inherited runtime locks; copies all state needed by the codec. - OffsetBufferAllocatorSnapshot CaptureSnapshot() const; - std::unique_ptr allocate(size_t size) override; void deallocate(AllocatedBuffer* handle) override; @@ -420,10 +387,6 @@ class OffsetBufferAllocator } private: - OffsetBufferAllocator( - OffsetBufferAllocatorSnapshot snapshot, - std::shared_ptr offset_allocator); - // metadata const std::string segment_name_; const size_t base_; diff --git a/mooncake-store/include/ha/snapshot/allocator_snapshot_codec.h b/mooncake-store/include/ha/snapshot/allocator_snapshot_codec.h deleted file mode 100644 index ae5e298da9..0000000000 --- a/mooncake-store/include/ha/snapshot/allocator_snapshot_codec.h +++ /dev/null @@ -1,22 +0,0 @@ -#pragma once - -#include -#include - -#include "allocator.h" -#include "serialize/serializer.h" - -namespace mooncake::ha { - -// Owns the snapshot representation of concrete physical allocators. Driver and -// SegmentPool lifecycle code remain unaware of serialization details. -class AllocatorSnapshotCodec final { - public: - static tl::expected Encode( - const OffsetBufferAllocatorSnapshot& snapshot, MsgpackPacker& packer); - - static tl::expected - Decode(const msgpack::object& object); -}; - -} // namespace mooncake::ha diff --git a/mooncake-store/include/ha/snapshot/local_ssd_codec.h b/mooncake-store/include/ha/snapshot/local_ssd_codec.h deleted file mode 100644 index 70336d9e2f..0000000000 --- a/mooncake-store/include/ha/snapshot/local_ssd_codec.h +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include -#include - -#include "local_ssd/persisted_state.h" -#include "serialize/serializer.h" - -namespace mooncake::ha { - -class LocalSsdCodec { - public: - static tl::expected Encode( - const LocalSsdPersistedState& state, MsgpackPacker& packer); - - static tl::expected Decode( - const msgpack::object* ld_value); -}; - -} // namespace mooncake::ha diff --git a/mooncake-store/include/ha/snapshot/store_resource_snapshot_codec.h b/mooncake-store/include/ha/snapshot/store_resource_snapshot_codec.h deleted file mode 100644 index af491d5931..0000000000 --- a/mooncake-store/include/ha/snapshot/store_resource_snapshot_codec.h +++ /dev/null @@ -1,39 +0,0 @@ -#pragma once - -#include -#include - -#include - -#include "local_ssd/persisted_state.h" -#include "segment/snapshot.h" -#include "types.h" - -namespace mooncake { -struct SerializationError; -} // namespace mooncake - -namespace mooncake::ha { - -struct StoreResourceSnapshot { - SegmentPoolSnapshot segment_pool; - LocalSsdPersistedState local_ssd; -}; - -// Codec for the existing ma/an/ms/cs/ld store-resource snapshot payload. -// Runtime components remain unaware of serialization. -class StoreResourceSnapshotCodec final { - public: - // Reads only owned snapshot data, independent of the source pool's - // lifetime. Preserves SegmentSerializer's wire format (OFFSET host memory, - // not CXL). - static tl::expected, SerializationError> Encode( - const SegmentPoolSnapshot& snapshot, - const LocalSsdPersistedState& local_ssd_state); - - // Decodes detached state without publishing resources or updating metrics. - static tl::expected Decode( - const std::vector& data); -}; - -} // namespace mooncake::ha diff --git a/mooncake-store/include/local_ssd/manager.h b/mooncake-store/include/local_ssd/manager.h index 2a55075850..df6a5576f8 100644 --- a/mooncake-store/include/local_ssd/manager.h +++ b/mooncake-store/include/local_ssd/manager.h @@ -11,7 +11,6 @@ #include #include -#include "local_ssd/persisted_state.h" #include "mutex.h" #include "tenant_id.h" #include "types.h" @@ -106,8 +105,6 @@ class LocalSsdManager { void RequestRemoveAll(const std::vector& clients); tl::expected ConsumeRemoveAll(const UUID& client_id); - LocalSsdPersistedState ExportPersistedState() const; - void RestorePersistedState(LocalSsdPersistedState state); void Clear(); private: diff --git a/mooncake-store/include/local_ssd/persisted_state.h b/mooncake-store/include/local_ssd/persisted_state.h deleted file mode 100644 index 861e568a26..0000000000 --- a/mooncake-store/include/local_ssd/persisted_state.h +++ /dev/null @@ -1,22 +0,0 @@ -#pragma once - -#include -#include - -#include "types.h" - -namespace mooncake { - -struct LocalSsdPersistedClient { - bool enable_offloading{false}; - int64_t total_capacity_bytes{0}; - - // The codec iterates this map directly, so ordering by the legacy encoded - // key makes the snapshot bytes deterministic without rewriting that key. - std::map pending_offloads; -}; - -// UUID ordering likewise makes client encoding deterministic. -using LocalSsdPersistedState = std::map; - -} // namespace mooncake diff --git a/mooncake-store/include/offset_allocator/offset_allocator.h b/mooncake-store/include/offset_allocator/offset_allocator.h index 1225bf782e..b2c9831033 100644 --- a/mooncake-store/include/offset_allocator/offset_allocator.h +++ b/mooncake-store/include/offset_allocator/offset_allocator.h @@ -13,7 +13,6 @@ #include #include "mutex.h" -#include "serialize/serializer.h" namespace mooncake::offset_allocator { typedef unsigned char uint8; @@ -24,7 +23,6 @@ using NodeIndex = uint32; // Forward declarations class OffsetAllocator; class __Allocator; -struct OffsetAllocatorSnapshot; static constexpr uint32 NUM_TOP_BINS = 32; static constexpr uint32 BINS_PER_LEAF = 8; @@ -48,7 +46,6 @@ struct OffsetAllocation { bool isNoSpace() const { return offset == NO_SPACE; } friend class __Allocator; - friend class Serializer; friend class OffsetAllocator; // for createHandleAtNode during recovery }; @@ -113,7 +110,6 @@ class OffsetAllocationHandle { uint64_t requested_size; friend class OffsetAllocatorTest; // for unit tests - friend class Serializer; }; struct OffsetAllocatorMetrics { @@ -195,15 +191,6 @@ class OffsetAllocator : public std::enable_shared_from_this { [[nodiscard]] OffsetAllocatorMetrics get_metrics() const; - // Deep-copy persisted layout without locking. Only for a forked snapshot - // child or externally quiesced state; the result owns no live resources. - OffsetAllocatorSnapshot CaptureSnapshot() const; - - // Consumes a captured snapshot and rebuilds the allocator from its layout. - // Validates before installing any state; returns a diagnostic on failure. - static tl::expected, std::string> Restore( - OffsetAllocatorSnapshot snapshot); - // Serialize the allocator with serializer. template void serialize_to(T& serializer) const; @@ -213,7 +200,6 @@ class OffsetAllocator : public std::enable_shared_from_this { private: friend class OffsetAllocationHandle; - friend class Serializer; // Internal method for Handle to free allocation (thread-safe) void freeAllocation(const OffsetAllocation& allocation, uint64_t size); @@ -242,12 +228,6 @@ class OffsetAllocator : public std::enable_shared_from_this { OffsetAllocator(uint64_t base, size_t size, uint32 init_capacity, uint32 max_capacity); - // Private constructor for creating from saved state with pre-constructed - // allocator This constructor is used during deserialization when we already - // have a reconstructed - OffsetAllocator(uint64_t base, size_t size, uint64_t multiplier_bits, - std::unique_ptr<__Allocator> allocator); - // Private constructor - initialize from serialized data template OffsetAllocator(T& serializer); @@ -281,9 +261,6 @@ class __Allocator { void serialize_to(T& serializer) const; private: - // Detaches the bin/node state for OffsetAllocator::CaptureSnapshot(). - __Allocator(const __Allocator&) = default; - uint32 insertNodeIntoBin(uint32 size, uint32 dataOffset); void removeNodeFromBin(uint32 nodeIndex); @@ -313,59 +290,9 @@ class __Allocator { uint32 m_freeOffset; friend class OffsetAllocatorTest; // for unit tests - friend class mooncake::Serializer<__Allocator>; - friend struct OffsetAllocatorSnapshot; friend class OffsetAllocator; // for visit_used_nodes / createHandleAtNode }; -// Owned allocation metadata. No runtime allocator, handles or mutex are -// retained. The internal bin/node layout is deeply copied during capture. -struct OffsetAllocatorSnapshot { - uint64_t base; - uint64_t multiplier_bits; - uint64_t capacity; - uint64_t allocated_size; - uint64_t allocated_num; - std::unique_ptr<__Allocator> layout; - - // The host-memory allocator supports up to 64Mi metadata nodes. This is - // independent of the smaller limit used by the storage backend's byte - // serializer, and must allow snapshots of large memory segments. - static constexpr uint32_t kMaxNodes = 1u << 26; - - // Used by the wire decoder before allocating/decompressing node storage. - static tl::expected ValidateNodeCapacity( - uint32_t size, uint32_t current_capacity, uint32_t max_capacity, - uint32_t free_offset); - - // Pure validation: no runtime allocator, metrics, logging or mutation. - // Checks headers before shifts, then layout topology and usage accounting. - // Requested sizes are not persisted per node, so allocated_size must lie - // within the range permitted by bin/unit rounding, not equal occupied - // bytes. - [[nodiscard]] tl::expected Validate() const; - - private: - enum class NodeState : uint8_t { kUnclassified, kUsed, kFree, kSpare }; - struct LayoutUsage { - uint64_t used_nodes = 0; - uint64_t min_used_bytes = 0; - uint64_t max_used_bytes = 0; - }; - - tl::expected ValidateHeader() const; - tl::expected ValidateUsedNodes( - std::vector& state) const; - tl::expected ValidateBins( - std::vector& state) const; - tl::expected ValidateFreeStack( - std::vector& state, uint32_t live_nodes) const; - tl::expected ValidateNeighborChain( - const std::vector& state, uint32_t live_nodes) const; - tl::expected ValidateAccounting( - const LayoutUsage& usage) const; -}; - // Template method implementations template void OffsetAllocator::serialize_to(T& serializer) const { diff --git a/mooncake-store/include/replica.h b/mooncake-store/include/replica.h index f35b82fbd4..8a3505a094 100644 --- a/mooncake-store/include/replica.h +++ b/mooncake-store/include/replica.h @@ -861,8 +861,7 @@ class Replica { data_; ReplicaStatus status_{ReplicaStatus::UNDEFINED}; - friend class Serializer; - friend class MasterService; // For MetadataSerializer to access next_id_ + friend class MasterService; // Standby restore advances next_id_ std::atomic refcnt_{0}; }; diff --git a/mooncake-store/include/segment.h b/mooncake-store/include/segment.h index 1fa97b109d..cec1cc40e7 100644 --- a/mooncake-store/include/segment.h +++ b/mooncake-store/include/segment.h @@ -14,7 +14,6 @@ #include "allocation_strategy.h" #include "allocator.h" -#include "local_ssd/persisted_state.h" #include "rpc_types.h" #include "segment/status.h" #include "segment/usage.h" @@ -336,47 +335,6 @@ class ScopedAllocatorAccess { std::shared_lock lock_; }; -/** - * @brief Lock-free data view of SegmentManager for safely accessing inherited - * data in forked child processes This class provides read-only access to - * SegmentManager private members without any lock operations - */ -class SegmentView { - public: - /** - * @brief Constructor - * @param segment_manager Reference to SegmentManager - */ - explicit SegmentView(const SegmentManager* segment_manager) - : segment_manager_(segment_manager) {} - - ErrorCode GetSegment(const std::shared_ptr& allocator, - Segment& segment) const; - - ErrorCode GetMountedSegment(const UUID& segment_id, - MountedSegment& mountedSegment) const; - - private: - const SegmentManager* segment_manager_; -}; - -class SegmentSerializer { - public: - explicit SegmentSerializer(SegmentManager* segment_manager) - : segment_manager_(segment_manager) {} - - tl::expected, SerializationError> Serialize( - const LocalSsdPersistedState& local_ssd_state); - - tl::expected Deserialize( - const std::vector& data); - - void Reset(); - - private: - SegmentManager* segment_manager_; -}; - class SegmentManager { public: /** @@ -417,8 +375,6 @@ class SegmentManager { client_by_name_, segment_mutex_); } - SegmentView getView() const { return SegmentView(this); } - /** * @brief Return current DRAM usage derived from mounted allocators. * @@ -473,9 +429,7 @@ class SegmentManager { HostSegmentIndex segments_by_host_; // host_id -> segment name -> segment // ids for allocatable segments friend class ScopedSegmentAccess; - friend class SegmentTest; // for unit tests - friend class SegmentView; // for fork serialize - friend class SegmentSerializer; // for fork serialize + friend class SegmentTest; // for unit tests }; class NoFSegmentManager { diff --git a/mooncake-store/include/segment/pool.h b/mooncake-store/include/segment/pool.h index a4a15aa339..1352c78452 100644 --- a/mooncake-store/include/segment/pool.h +++ b/mooncake-store/include/segment/pool.h @@ -8,7 +8,6 @@ #include "placement/replica_allocator.h" #include "segment/catalog.h" #include "segment/region_driver.h" -#include "segment/snapshot.h" #include "segment/usage.h" #include "storage_usage.h" @@ -30,18 +29,6 @@ class SegmentPool final { WriteAccess AcquireWriteAccess(); ReadAccess AcquireReadAccess() const; - // Capture detached data without acquiring pool/allocator locks. Only call - // in a forked snapshot child or with all relevant state externally - // quiesced. After capture, the result is independent of this pool's - // lifetime/mutations. - tl::expected CaptureSnapshot() const; - - // Consumes the snapshot, staging all resources before replacing the pool. - // Preparation failure leaves published state unchanged. Temporary readers - // may omit capacity accounting. - tl::expected RestoreSnapshot( - SegmentPoolSnapshot snapshot, bool account_capacity_metrics); - // Holds the Pool read lock through candidate selection and allocation. template tl::expected, ErrorCode> AllocateReplicas( diff --git a/mooncake-store/include/segment/pool_write_access.h b/mooncake-store/include/segment/pool_write_access.h index 4cd1fabb2c..36a4cc4f72 100644 --- a/mooncake-store/include/segment/pool_write_access.h +++ b/mooncake-store/include/segment/pool_write_access.h @@ -24,8 +24,6 @@ class SegmentPool::WriteAccess final { class RegionUnmountTxn; class RegionGracefulUnmountTxn; - enum class AdoptMode { Insert, ReplacePool }; - WriteAccess(AccessKey, SegmentPool& segment_pool); ErrorCode MountSegment(const Segment& segment, const UUID& client_id); @@ -36,11 +34,9 @@ class SegmentPool::WriteAccess final { tl::expected PrepareRestore( const Segment& segment, const UUID& client_id, std::span descriptors); - // ReplacePool ignores conflicts with the old catalog. The caller must - // validate the replacement and Clear() before committing transactions. tl::expected PrepareAdopt( MountedRegion mounted, std::shared_ptr allocator, - bool account_capacity_metrics, AdoptMode mode = AdoptMode::Insert); + bool account_capacity_metrics); void Clear() noexcept; tl::expected PrepareUnmount( diff --git a/mooncake-store/include/segment/snapshot.h b/mooncake-store/include/segment/snapshot.h deleted file mode 100644 index 32040e8297..0000000000 --- a/mooncake-store/include/segment/snapshot.h +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include -#include - -#include "allocator.h" -#include "segment/mounted_region.h" - -namespace mooncake { - -struct RegionSnapshot { - MountedRegion mounted; - OffsetBufferAllocatorSnapshot allocator; -}; - -// Detached, owning snapshot of the resources in the legacy host-memory format. -// All metadata, including allocator bins/nodes, is copied. Encoding and -// destroying this value never accesses a SegmentPool or changes runtime -// capacity/usage. -struct SegmentPoolSnapshot { - BufferAllocatorType memory_allocator_type = BufferAllocatorType::OFFSET; - std::vector active_names; - std::vector regions; -}; - -} // namespace mooncake diff --git a/mooncake-store/include/segment/status.h b/mooncake-store/include/segment/status.h index a7f41bfdef..1c2ab3c5d7 100644 --- a/mooncake-store/include/segment/status.h +++ b/mooncake-store/include/segment/status.h @@ -13,20 +13,6 @@ enum class SegmentStatus { UNMOUNTING, }; -// Rejects raw integer values that do not map to a declared status. -constexpr bool IsKnownSegmentStatus(SegmentStatus status) noexcept { - switch (status) { - case SegmentStatus::UNDEFINED: - case SegmentStatus::OK: - case SegmentStatus::DRAINING: - case SegmentStatus::DRAINED: - case SegmentStatus::GRACEFULLY_UNMOUNTING: - case SegmentStatus::UNMOUNTING: - return true; - } - return false; -} - constexpr int SegmentStatusAvailabilityRank(SegmentStatus status) noexcept { switch (status) { case SegmentStatus::OK: diff --git a/mooncake-store/include/segment_allocator_registration.h b/mooncake-store/include/segment_allocator_registration.h index fbdca4c318..f052c81fdd 100644 --- a/mooncake-store/include/segment_allocator_registration.h +++ b/mooncake-store/include/segment_allocator_registration.h @@ -8,9 +8,6 @@ namespace mooncake { class ScopedNoFSegmentAccess; class ScopedSegmentAccess; -class SegmentSerializer; -template -class Serializer; class SegmentAllocatorRegistration { public: [[nodiscard]] bool IsServing() const; @@ -35,8 +32,6 @@ class SegmentAllocatorRegistration { friend class AllocatorManager; friend class ScopedNoFSegmentAccess; friend class ScopedSegmentAccess; - friend class SegmentSerializer; - friend class Serializer; }; } // namespace mooncake diff --git a/mooncake-store/include/serialize/serializer.h b/mooncake-store/include/serialize/serializer.h deleted file mode 100644 index 88ab09e266..0000000000 --- a/mooncake-store/include/serialize/serializer.h +++ /dev/null @@ -1,150 +0,0 @@ -#pragma once - -#include -#include - -#include -#include -#include - -namespace mooncake::offset_allocator { -class __Allocator; -class OffsetAllocator; -class OffsetAllocationHandle; -} // namespace mooncake::offset_allocator - -namespace mooncake { -class AllocatedBuffer; -class Replica; -class MasterService; -class BufferAllocatorBase; -class SegmentView; -enum class ErrorCode; -struct SerializationError; -class MountedSegment; -class OffsetBufferAllocator; - -using MsgpackPacker = msgpack::packer; - -// Generic serialization interface -template -class Serializer; - -// Serializer specialization for __Allocator -template <> -class Serializer { - public: - using PointerType = std::unique_ptr; - - // Layout validation lives with the snapshot so every restore path shares - // one checker (OffsetAllocatorSnapshot::Validate); this codec only bounds - // the scalar capacities before decompressing node payloads. - static tl::expected serialize( - const offset_allocator::__Allocator &allocator, MsgpackPacker &packer); - - static tl::expected deserialize( - const msgpack::object &obj); -}; - -// Serializer specialization for OffsetAllocator -template <> -class Serializer { - public: - using PointerType = std::shared_ptr; - - static tl::expected serialize( - const offset_allocator::OffsetAllocator &allocator, - MsgpackPacker &packer); - - static tl::expected deserialize( - const msgpack::object &obj); -}; - -// Serializer specialization for OffsetAllocationHandle -template <> -class Serializer { - public: - using PointerType = - std::shared_ptr; - - static tl::expected serialize( - const offset_allocator::OffsetAllocationHandle &handle, - MsgpackPacker &packer); - - static tl::expected deserialize( - const msgpack::object &obj, - const std::shared_ptr &allocator); -}; - -// Serializer specialization for AllocatedBuffer -template <> -class Serializer { - public: - using PointerType = std::unique_ptr; - - static tl::expected serialize( - const AllocatedBuffer &buffer, const SegmentView &segment_view, - MsgpackPacker &packer); - - static tl::expected deserialize( - const msgpack::object &obj, const SegmentView &segment_view); -}; - -// Serializer specialization for Replica (interface declaration) -template <> -class Serializer { - public: - using PointerType = std::shared_ptr; - - static tl::expected serialize( - const Replica &replica, const SegmentView &segment_view, - MsgpackPacker &packer); - - static tl::expected deserialize( - const msgpack::object &obj, const SegmentView &segment_view); -}; - -template <> -class Serializer { - public: - static tl::expected serialize( - const MountedSegment &mounted_segment, MsgpackPacker &packer); - - static tl::expected deserialize( - const msgpack::object &obj); -}; - -// Serializer specialization for OffsetBufferAllocator (interface declaration) -template <> -class Serializer { - public: - using PointerType = std::shared_ptr; - - static tl::expected serialize( - const OffsetBufferAllocator &allocator, MsgpackPacker &packer); - - static tl::expected deserialize( - const msgpack::object &obj); -}; - -// Generic serialization helper class -class SerializationHelper { - public: - // Serialize uint32_t (little-endian) - static void serializeUint32(uint32_t value, std::vector &out) { - out.push_back(static_cast(value & 0xFF)); - out.push_back(static_cast((value >> 8) & 0xFF)); - out.push_back(static_cast((value >> 16) & 0xFF)); - out.push_back(static_cast((value >> 24) & 0xFF)); - } - - // Deserialize uint32_t (little-endian) - static uint32_t deserializeUint32(const uint8_t *data) { - return static_cast(data[0]) | - (static_cast(data[1]) << 8) | - (static_cast(data[2]) << 16) | - (static_cast(data[3]) << 24); - } -}; - -} // namespace mooncake diff --git a/mooncake-store/include/task_manager.h b/mooncake-store/include/task_manager.h index b5cf8c69ba..f2cf596997 100644 --- a/mooncake-store/include/task_manager.h +++ b/mooncake-store/include/task_manager.h @@ -181,13 +181,9 @@ class ScopedTaskWriteAccess { ErrorCode complete_task(const UUID& client_id, const UUID& task_id, TaskStatus status, const std::string& message); - void restore_task(Task&& task); - ErrorCode fail_task_if_pending(const UUID& task_id, const std::string& message); - void clear_all(); - void prune_finished_tasks(); void prune_expired_tasks(); @@ -259,22 +255,4 @@ class ClientTaskManager { std::deque finished_task_history_ GUARDED_BY(mutex_); }; -class TaskManagerSerializer { - static constexpr size_t kTaskSerializedFields = 8; - static constexpr size_t kMaxDecompressedSize = 1024 * 1024 * 1024; // 1 GB - - public: - explicit TaskManagerSerializer(ClientTaskManager* task_manager) - : task_manager_(task_manager) {} - - tl::expected, SerializationError> Serialize(); - - tl::expected Deserialize( - const std::vector& data); - - void Reset(); - - private: - ClientTaskManager* task_manager_; -}; } // namespace mooncake diff --git a/mooncake-store/include/types.h b/mooncake-store/include/types.h index d26a6126f6..713058dbb0 100644 --- a/mooncake-store/include/types.h +++ b/mooncake-store/include/types.h @@ -428,17 +428,6 @@ inline std::ostream& operator<<(std::ostream& os, return os << toString(errorCode); } -struct SerializationError { - ErrorCode code; - std::string message; - - SerializationError(ErrorCode c, const std::string& msg) - : code(c), message(msg) {} - - // Construct from ErrorCode - explicit SerializationError(ErrorCode c) : code(c), message(toString(c)) {} -}; - /** * @brief Represents a contiguous memory region */ diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index 91fa180b66..788a0a633b 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -41,9 +41,6 @@ set(MOONCAKE_STORE_MASTER_SOURCES segment/catalog.cpp segment/pool.cpp segment/pool_mount.cpp - segment/snapshot.cpp - ha/snapshot/allocator_snapshot_codec.cpp - ha/snapshot/store_resource_snapshot_codec.cpp segment/pool_write_access.cpp segment/pool_read_access.cpp segment.cpp @@ -57,7 +54,6 @@ set(MOONCAKE_STORE_MASTER_SOURCES offset_allocator.cpp http_metadata_server.cpp embedded_master.cpp - serialize/serializer.cpp storage/distributed/shard_allocator.cpp storage/distributed/immutable_bucket_allocator.cpp ha/leadership/leader_coordinator_factory.cpp @@ -80,7 +76,6 @@ set(MOONCAKE_STORE_MASTER_SOURCES ha/snapshot/batch_oplog/batch_oplog_snapshot_gc.cpp ha/snapshot/batch_oplog/batch_oplog_pruning_coordinator.cpp ha/snapshot/snapshot_maintenance_lease.cpp - ha/snapshot/local_ssd_codec.cpp ha/snapshot/object/snapshot_object_store.cpp ha/snapshot/object/backends/local/local_file_snapshot_object_store.cpp ha/snapshot/object/backends/s3/s3_snapshot_object_store.cpp diff --git a/mooncake-store/src/allocator.cpp b/mooncake-store/src/allocator.cpp index 58ab2b1e44..2ff2969c4b 100644 --- a/mooncake-store/src/allocator.cpp +++ b/mooncake-store/src/allocator.cpp @@ -373,9 +373,7 @@ OffsetBufferAllocator::OffsetBufferAllocator(std::string segment_name, max_capacity = std::max(max_capacity, static_cast(1024 * 1024)); max_capacity = - std::min(max_capacity, - static_cast( - offset_allocator::OffsetAllocatorSnapshot::kMaxNodes)); + std::min(max_capacity, static_cast(64 * 1024 * 1024)); // Create the offset allocator offset_allocator_ = offset_allocator::OffsetAllocator::create( base, size, static_cast(init_capacity), @@ -393,79 +391,6 @@ OffsetBufferAllocator::OffsetBufferAllocator(std::string segment_name, } } -OffsetBufferAllocator::OffsetBufferAllocator( - OffsetBufferAllocatorSnapshot snapshot, - std::shared_ptr offset_allocator) - : segment_name_(std::move(snapshot.segment_name)), - base_(snapshot.base), - total_size_(snapshot.capacity), - transport_endpoint_(std::move(snapshot.transport_endpoint)), - replica_type_(ReplicaType::MEMORY), - offset_allocator_(std::move(offset_allocator)) {} - -tl::expected -OffsetBufferAllocatorSnapshot::ValidateMetadata() const { - if (allocation_state.base != base || - allocation_state.capacity != capacity) { - return tl::make_unexpected(fmt::format( - "allocator bounds mismatch: outer base={}, capacity={}; inner " - "base={}, capacity={}", - base, capacity, allocation_state.base, allocation_state.capacity)); - } - if (used_bytes > capacity || - used_bytes > - static_cast(std::numeric_limits::max())) { - return tl::make_unexpected( - fmt::format("invalid usage {} for capacity {} or signed metrics", - used_bytes, capacity)); - } - if (used_bytes != allocation_state.allocated_size) { - return tl::make_unexpected( - fmt::format("used_bytes {} does not match allocated_size {}", - used_bytes, allocation_state.allocated_size)); - } - return {}; -} - -tl::expected OffsetBufferAllocatorSnapshot::Validate() - const { - if (auto valid = ValidateMetadata(); !valid) return valid; - return allocation_state.Validate(); -} - -tl::expected, ErrorCode> -OffsetBufferAllocator::Restore(OffsetBufferAllocatorSnapshot snapshot) { - const auto used_bytes = snapshot.used_bytes; - if (auto valid = snapshot.ValidateMetadata(); !valid) { - LOG(ERROR) << "OffsetBufferAllocator::Restore rejected snapshot for " - << snapshot.segment_name << ": " << valid.error(); - return tl::make_unexpected(ErrorCode::INVALID_PARAMS); - } - // The inner factory validates its complete snapshot and logs the specific - // failure before constructing an allocator. No metrics are published yet. - auto offset_allocator = offset_allocator::OffsetAllocator::Restore( - std::move(snapshot.allocation_state)); - if (!offset_allocator) { - return tl::make_unexpected(ErrorCode::INVALID_PARAMS); - } - - auto allocator = - std::shared_ptr(new OffsetBufferAllocator( - std::move(snapshot), std::move(*offset_allocator))); - // No allocation handles exist yet. Pair restored usage with the same - // accounting that allocate()/destruction use, before publishing the object. - MasterMetricManager::instance().inc_allocated_mem_size( - allocator->segment_name_, static_cast(used_bytes)); - allocator->SetUsageBytesForRestore(used_bytes); - return allocator; -} - -OffsetBufferAllocatorSnapshot OffsetBufferAllocator::CaptureSnapshot() const { - return {segment_name_, base_, - total_size_, size(), - transport_endpoint_, offset_allocator_->CaptureSnapshot()}; -} - OffsetBufferAllocator::~OffsetBufferAllocator() { if (replica_type_ == ReplicaType::MEMORY) { MasterMetricManager::instance().dec_allocated_mem_size(segment_name_, diff --git a/mooncake-store/src/ha/snapshot/allocator_snapshot_codec.cpp b/mooncake-store/src/ha/snapshot/allocator_snapshot_codec.cpp deleted file mode 100644 index f79640cf5e..0000000000 --- a/mooncake-store/src/ha/snapshot/allocator_snapshot_codec.cpp +++ /dev/null @@ -1,97 +0,0 @@ -#include "ha/snapshot/allocator_snapshot_codec.h" - -#include - -#include - -namespace mooncake::ha { - -tl::expected AllocatorSnapshotCodec::Encode( - const OffsetBufferAllocatorSnapshot& snapshot, MsgpackPacker& packer) { - if (!snapshot.allocation_state.layout) { - return tl::make_unexpected(SerializationError( - ErrorCode::SERIALIZE_FAIL, "snapshot allocator has no layout")); - } - // Preserve the legacy OffsetBufferAllocator and OffsetAllocator arrays. - packer.pack_array(6); - packer.pack(snapshot.segment_name); - packer.pack(static_cast(snapshot.base)); - packer.pack(static_cast(snapshot.capacity)); - packer.pack(static_cast(snapshot.used_bytes)); - packer.pack(snapshot.transport_endpoint); - const auto& layout = snapshot.allocation_state; - packer.pack_array(6); - packer.pack(layout.base); - packer.pack(layout.multiplier_bits); - packer.pack(layout.capacity); - packer.pack(layout.allocated_size); - packer.pack(layout.allocated_num); - return Serializer::serialize(*layout.layout, - packer); -} - -tl::expected -AllocatorSnapshotCodec::Decode(const msgpack::object& object) { - if (object.type != msgpack::type::ARRAY || object.via.array.size != 6) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "snapshot OffsetBufferAllocator is not an array[6]")); - } - - try { - const auto* array = object.via.array.ptr; - std::string segment_name = array[0].as(); - auto base = static_cast(array[1].as()); - auto total_size = static_cast(array[2].as()); - auto used_bytes = static_cast(array[3].as()); - std::string transport_endpoint = array[4].as(); - - if (array[5].type != msgpack::type::ARRAY || - array[5].via.array.size != 6) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "snapshot OffsetAllocator is not an array[6]")); - } - const auto* state = array[5].via.array.ptr; - auto layout = - Serializer::deserialize(state[5]); - if (!layout) { - return tl::make_unexpected(layout.error()); - } - - OffsetBufferAllocatorSnapshot snapshot{ - .segment_name = std::move(segment_name), - .base = base, - .capacity = total_size, - .used_bytes = used_bytes, - .transport_endpoint = std::move(transport_endpoint), - .allocation_state = - offset_allocator::OffsetAllocatorSnapshot{ - .base = state[0].as(), - .multiplier_bits = state[1].as(), - .capacity = state[2].as(), - .allocated_size = state[3].as(), - .allocated_num = state[4].as(), - .layout = std::move(*layout), - }, - }; - // The snapshot owns both the outer metadata and the inner layout, so - // one validation pass covers the bounds/base/usage consistency and the - // allocator topology. Decode propagates diagnostics to its caller; - // direct restore factories also log validation failures. - if (auto valid = snapshot.Validate(); !valid) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("snapshot OffsetBufferAllocator {}", - valid.error()))); - } - return snapshot; - } catch (const std::exception& error) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("decode OffsetBufferAllocator failed: {}", - error.what()))); - } -} - -} // namespace mooncake::ha diff --git a/mooncake-store/src/ha/snapshot/local_ssd_codec.cpp b/mooncake-store/src/ha/snapshot/local_ssd_codec.cpp deleted file mode 100644 index 0f9d3d63c3..0000000000 --- a/mooncake-store/src/ha/snapshot/local_ssd_codec.cpp +++ /dev/null @@ -1,153 +0,0 @@ -#include "ha/snapshot/local_ssd_codec.h" - -#include - -#include "tenant_id.h" - -namespace mooncake::ha { -namespace { - -bool IsMsgpackInteger(const msgpack::object& object) { - return object.type == msgpack::type::POSITIVE_INTEGER || - object.type == msgpack::type::NEGATIVE_INTEGER; -} - -SerializationError DecodeError(std::string message) { - return SerializationError(ErrorCode::DESERIALIZE_FAIL, message); -} - -} // namespace - -tl::expected LocalSsdCodec::Encode( - const LocalSsdPersistedState& state, MsgpackPacker& packer) { - packer.pack_map(state.size()); - for (const auto& [client_id, client] : state) { - packer.pack(UuidToString(client_id)); - packer.pack_array(2 + client.pending_offloads.size() * 2 + 1); - packer.pack(client.enable_offloading); - packer.pack(static_cast(client.pending_offloads.size())); - for (const auto& [encoded_key, task] : client.pending_offloads) { - packer.pack(encoded_key); - packer.pack_array(3); - packer.pack(task.tenant_id); - packer.pack(task.key); - packer.pack(task.size); - } - packer.pack(client.total_capacity_bytes); - } - return {}; -} - -tl::expected LocalSsdCodec::Decode( - const msgpack::object* ld_value) { - LocalSsdPersistedState state; - if (ld_value == nullptr) { - return state; - } - if (ld_value->type != msgpack::type::MAP) { - return tl::unexpected( - DecodeError("deserialize local_disk_segments is not map")); - } - - try { - for (uint32_t i = 0; i < ld_value->via.map.size; ++i) { - const auto& client_key = ld_value->via.map.ptr[i].key; - const auto& client_value = ld_value->via.map.ptr[i].val; - if (client_key.type != msgpack::type::STR) { - return tl::unexpected(DecodeError( - "deserialize local_disk_segments client key is not " - "string")); - } - std::string client_uuid(client_key.via.str.ptr, - client_key.via.str.size); - UUID client_id; - if (!StringToUuid(client_uuid, client_id)) { - return tl::unexpected(DecodeError(fmt::format( - "deserialize local_disk_segments client uuid {} is " - "invalid", - client_uuid))); - } - if (client_value.type != msgpack::type::ARRAY || - client_value.via.array.size < 2) { - return tl::unexpected(DecodeError( - "deserialize local_disk_segments value is not valid " - "array")); - } - if (client_value.via.array.ptr[0].type != msgpack::type::BOOLEAN || - !IsMsgpackInteger(client_value.via.array.ptr[1])) { - return tl::unexpected(DecodeError( - "deserialize local_disk_segments header is invalid")); - } - - LocalSsdPersistedClient client; - client.enable_offloading = client_value.via.array.ptr[0].as(); - uint64_t count = client_value.via.array.ptr[1].as(); - for (uint64_t task_number = 0; task_number < count; ++task_number) { - size_t key_index = 2 + task_number * 2; - size_t task_index = key_index + 1; - if (task_index >= client_value.via.array.size) { - return tl::unexpected(DecodeError( - "deserialize local_disk_segments offloading_objects " - "out of bounds")); - } - const auto& key_object = client_value.via.array.ptr[key_index]; - if (key_object.type != msgpack::type::STR) { - return tl::unexpected(DecodeError( - "deserialize local_disk_segments offloading key is " - "not string")); - } - std::string encoded_key(key_object.via.str.ptr, - key_object.via.str.size); - const auto& task_object = - client_value.via.array.ptr[task_index]; - OffloadTaskItem task; - if (task_object.type == msgpack::type::ARRAY && - task_object.via.array.size == 3) { - if (task_object.via.array.ptr[0].type != - msgpack::type::STR || - task_object.via.array.ptr[1].type != - msgpack::type::STR || - !IsMsgpackInteger(task_object.via.array.ptr[2])) { - return tl::unexpected(DecodeError( - "deserialize local_disk_segments offloading task " - "is invalid")); - } - task.tenant_id = - task_object.via.array.ptr[0].as(); - task.key = task_object.via.array.ptr[1].as(); - task.size = task_object.via.array.ptr[2].as(); - } else { - if (!IsMsgpackInteger(task_object)) { - return tl::unexpected(DecodeError( - "deserialize local_disk_segments legacy " - "offloading size is not integer")); - } - auto [tenant_id, object_key] = - TenantId::ParseScopedKey(encoded_key); - task = OffloadTaskItem{.tenant_id = tenant_id.value(), - .key = std::move(object_key), - .size = task_object.as()}; - } - client.pending_offloads.insert_or_assign(std::move(encoded_key), - std::move(task)); - } - - const size_t capacity_index = 2 + count * 2; - if (client_value.via.array.size > capacity_index && - IsMsgpackInteger(client_value.via.array.ptr[capacity_index])) { - client.total_capacity_bytes = - client_value.via.array.ptr[capacity_index].as(); - } - state.insert_or_assign(client_id, std::move(client)); - } - } catch (const msgpack::type_error& error) { - return tl::unexpected(DecodeError(fmt::format( - "deserialize local_disk_segments type error: {}", error.what()))); - } catch (const std::exception& error) { - return tl::unexpected(DecodeError(fmt::format( - "deserialize local_disk_segments failed: {}", error.what()))); - } - return state; -} - -} // namespace mooncake::ha diff --git a/mooncake-store/src/ha/snapshot/store_resource_snapshot_codec.cpp b/mooncake-store/src/ha/snapshot/store_resource_snapshot_codec.cpp deleted file mode 100644 index 1ff6087ca7..0000000000 --- a/mooncake-store/src/ha/snapshot/store_resource_snapshot_codec.cpp +++ /dev/null @@ -1,388 +0,0 @@ -#include "ha/snapshot/store_resource_snapshot_codec.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include - -#include "ha/snapshot/allocator_snapshot_codec.h" -#include "ha/snapshot/local_ssd_codec.h" -#include "segment/snapshot.h" -#include "serialize/serializer.h" -#include "common/zstd_util.h" - -namespace mooncake::ha { -namespace { - -struct DecodedMountedRegion final { - MountedRegion mounted; - std::optional allocator; -}; - -tl::expected EncodeMountedRegion( - const MountedRegion& mounted, - const OffsetBufferAllocatorSnapshot& allocator, MsgpackPacker& packer) { - // Preserved wire shape: [segment_id, segment_name, segment_base, - // segment_size, te_endpoint, status, has_allocator, allocator, host_id]. - packer.pack_array(9); - packer.pack(UuidToString(mounted.segment.id)); - packer.pack(mounted.segment.name); - packer.pack(static_cast(mounted.segment.base)); - packer.pack(static_cast(mounted.segment.size)); - packer.pack(mounted.segment.te_endpoint); - packer.pack(static_cast(mounted.status)); - packer.pack(true); - auto encoded = AllocatorSnapshotCodec::Encode(allocator, packer); - if (!encoded) { - return tl::make_unexpected(encoded.error()); - } - packer.pack(mounted.segment.host_id); - return {}; -} - -tl::expected DecodeMountedRegion( - const msgpack::object& object) { - if (object.type != msgpack::type::ARRAY || object.via.array.size < 8) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "snapshot MountedRegion is not an array with at least 8 fields")); - } - - try { - const auto* array = object.via.array.ptr; - UUID id; - if (!StringToUuid(array[0].as(), id)) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot MountedRegion has invalid UUID")); - } - - DecodedMountedRegion decoded; - decoded.mounted.segment.id = id; - decoded.mounted.segment.name = array[1].as(); - decoded.mounted.segment.base = - static_cast(array[2].as()); - decoded.mounted.segment.size = - static_cast(array[3].as()); - decoded.mounted.segment.te_endpoint = array[4].as(); - decoded.mounted.status = - static_cast(array[5].as()); - if (!IsKnownSegmentStatus(decoded.mounted.status)) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "snapshot MountedRegion has invalid status")); - } - decoded.mounted.kind = RegionKind::HOST_MEMORY; - - if (array[6].as()) { - auto allocator = AllocatorSnapshotCodec::Decode(array[7]); - if (!allocator) { - return tl::make_unexpected(allocator.error()); - } - decoded.allocator = std::move(*allocator); - } - if (object.via.array.size >= 9) { - decoded.mounted.segment.host_id = array[8].as(); - } - return decoded; - } catch (const std::exception& error) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - std::string("decode MountedRegion failed: ") + error.what())); - } -} - -tl::expected, SerializationError> -IndexFields(const msgpack::object& object) { - if (object.type != msgpack::type::MAP) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot SegmentPool payload is not a map")); - } - std::map fields; - for (uint32_t i = 0; i < object.via.map.size; ++i) { - const auto& key = object.via.map.ptr[i].key; - if (key.type == msgpack::type::STR) { - fields.emplace(std::string(key.via.str.ptr, key.via.str.size), - &object.via.map.ptr[i].val); - } - } - return fields; -} - -} // namespace - -tl::expected, SerializationError> -StoreResourceSnapshotCodec::Encode( - const SegmentPoolSnapshot& snapshot, - const LocalSsdPersistedState& local_ssd_state) { - if (snapshot.memory_allocator_type != BufferAllocatorType::OFFSET) { - return tl::make_unexpected(SerializationError( - ErrorCode::SERIALIZE_UNSUPPORTED, - "snapshot SegmentPool memory driver is not offset")); - } - msgpack::sbuffer buffer; - MsgpackPacker packer(&buffer); - packer.pack_map(5); - - packer.pack("ma"); - packer.pack(static_cast(snapshot.memory_allocator_type)); - - packer.pack("an"); - packer.pack_array(snapshot.active_names.size()); - for (const auto& name : snapshot.active_names) { - packer.pack(name); - } - - packer.pack("ms"); - packer.pack_map(snapshot.regions.size()); - for (const auto& region : snapshot.regions) { - packer.pack(UuidToString(region.mounted.segment.id)); - auto encoded = - EncodeMountedRegion(region.mounted, region.allocator, packer); - if (!encoded) { - return tl::make_unexpected(encoded.error()); - } - } - - packer.pack("cs"); - std::map> clients; - for (const auto& region : snapshot.regions) { - clients[region.mounted.client_id].push_back(region.mounted.segment.id); - } - packer.pack_map(clients.size()); - for (const auto& [client, regions] : clients) { - packer.pack(UuidToString(client)); - packer.pack_array(regions.size()); - for (const auto& id : regions) { - packer.pack(UuidToString(id)); - } - } - - packer.pack("ld"); - auto local_ssd = LocalSsdCodec::Encode(local_ssd_state, packer); - if (!local_ssd) { - return tl::make_unexpected(local_ssd.error()); - } - return zstd_compress(reinterpret_cast(buffer.data()), - buffer.size(), 3); -} - -tl::expected -StoreResourceSnapshotCodec::Decode(const std::vector& data) { - std::vector decompressed; - try { - decompressed = zstd_decompress(data); - } catch (const std::exception& error) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - std::string("decompress SegmentPool snapshot failed: ") + - error.what())); - } - - msgpack::object_handle handle; - try { - handle = - msgpack::unpack(reinterpret_cast(decompressed.data()), - decompressed.size()); - } catch (const std::exception& error) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - std::string("unpack SegmentPool snapshot failed: ") + - error.what())); - } - auto fields = IndexFields(handle.get()); - if (!fields) { - return tl::make_unexpected(fields.error()); - } - - const auto local_ssd_field = fields->find("ld"); - auto local_ssd = LocalSsdCodec::Decode( - local_ssd_field == fields->end() ? nullptr : local_ssd_field->second); - if (!local_ssd) { - return tl::make_unexpected(local_ssd.error()); - } - - const auto allocator_field = fields->find("ma"); - if (allocator_field == fields->end() || - allocator_field->second->type != msgpack::type::POSITIVE_INTEGER || - allocator_field->second->via.u64 != - static_cast(BufferAllocatorType::OFFSET)) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot SegmentPool allocator type mismatch")); - } - - std::vector active_names; - const auto names_field = fields->find("an"); - if (names_field != fields->end()) { - if (names_field->second->type != msgpack::type::ARRAY) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot active names is not an array")); - } - for (uint32_t i = 0; i < names_field->second->via.array.size; ++i) { - const auto& name = names_field->second->via.array.ptr[i]; - if (name.type != msgpack::type::STR) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot active name is not a string")); - } - active_names.emplace_back(name.via.str.ptr, name.via.str.size); - } - } - - std::unordered_map> decoded; - const auto mounted_field = fields->find("ms"); - if (mounted_field != fields->end()) { - if (mounted_field->second->type != msgpack::type::MAP) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot mounted regions is not a map")); - } - for (uint32_t i = 0; i < mounted_field->second->via.map.size; ++i) { - const auto& item = mounted_field->second->via.map.ptr[i]; - if (item.key.type != msgpack::type::STR) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot region id is not a string")); - } - UUID id; - if (!StringToUuid( - std::string(item.key.via.str.ptr, item.key.via.str.size), - id)) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot region id is invalid")); - } - auto region = DecodeMountedRegion(item.val); - if (!region || region->mounted.segment.id != id || - !region->allocator) { - return tl::make_unexpected( - region ? SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "snapshot region id or allocator mismatch") - : region.error()); - } - const auto& segment = region->mounted.segment; - if (segment.name.empty() || segment.base == 0 || - segment.size == 0 || - segment.base > - std::numeric_limits::max() - segment.size || - region->allocator->segment_name != segment.name || - region->allocator->transport_endpoint != segment.te_endpoint || - region->allocator->base != segment.base || - region->allocator->capacity != segment.size) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "snapshot contains invalid adopted resource")); - } - if (!decoded.emplace(id, std::move(*region)).second) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot contains duplicate region")); - } - } - } - - std::set owned_regions; - const auto clients_field = fields->find("cs"); - if (clients_field != fields->end()) { - if (clients_field->second->type != msgpack::type::MAP) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot client regions is not a map")); - } - for (uint32_t i = 0; i < clients_field->second->via.map.size; ++i) { - const auto& item = clients_field->second->via.map.ptr[i]; - if (item.key.type != msgpack::type::STR || - item.val.type != msgpack::type::ARRAY) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "snapshot contains invalid client entry")); - } - UUID client; - if (!StringToUuid( - std::string(item.key.via.str.ptr, item.key.via.str.size), - client)) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot contains invalid client id")); - } - for (uint32_t j = 0; j < item.val.via.array.size; ++j) { - const auto& id_object = item.val.via.array.ptr[j]; - UUID id; - if (id_object.type != msgpack::type::STR || - !StringToUuid(std::string(id_object.via.str.ptr, - id_object.via.str.size), - id) || - !decoded.contains(id) || !owned_regions.insert(id).second) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "snapshot contains invalid region ownership")); - } - decoded.at(id).mounted.client_id = client; - } - } - } - if (owned_regions.size() != decoded.size()) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "snapshot region is missing an owner")); - } - - std::unordered_map owner_by_name; - std::set expected_active_names; - for (const auto& [_, region] : decoded) { - auto [owner, inserted] = owner_by_name.emplace( - region.mounted.segment.name, region.mounted.client_id); - if (!inserted && owner->second != region.mounted.client_id) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "snapshot logical region group has multiple owners")); - } - if (region.mounted.status == SegmentStatus::OK) { - expected_active_names.insert(region.mounted.segment.name); - } - } - - std::set decoded_active_names; - for (const auto& name : active_names) { - if (!decoded_active_names.insert(name).second || - !expected_active_names.contains(name)) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "snapshot contains invalid active region names")); - } - } - if (decoded_active_names != expected_active_names) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "snapshot active region names do not match mounted regions")); - } - - StoreResourceSnapshot snapshot; - snapshot.segment_pool.active_names = std::move(active_names); - snapshot.segment_pool.regions.reserve(decoded.size()); - for (auto& [id, region] : decoded) { - snapshot.segment_pool.regions.push_back( - {std::move(region.mounted), std::move(*region.allocator)}); - } - // Match CaptureSnapshot's canonical ID order for stable re-encoding. - std::sort(snapshot.segment_pool.regions.begin(), - snapshot.segment_pool.regions.end(), - [](const auto& lhs, const auto& rhs) { - return lhs.mounted.segment.id < rhs.mounted.segment.id; - }); - snapshot.local_ssd = std::move(*local_ssd); - return snapshot; -} - -} // namespace mooncake::ha diff --git a/mooncake-store/src/local_ssd/manager.cpp b/mooncake-store/src/local_ssd/manager.cpp index dda7b8de71..add47977c8 100644 --- a/mooncake-store/src/local_ssd/manager.cpp +++ b/mooncake-store/src/local_ssd/manager.cpp @@ -353,46 +353,6 @@ tl::expected LocalSsdManager::ConsumeRemoveAll( return client->record->mailbox.ConsumeRemoveAll(); } -LocalSsdPersistedState LocalSsdManager::ExportPersistedState() const { - auto clients = SnapshotClients(); - LocalSsdPersistedState state; - for (const auto& client_access : clients) { - const auto& record = client_access.record; - LocalSsdPersistedClient client; - { - MutexLocker stats_lock(&record->stats_mutex); - MutexLocker mailbox_lock(&record->mailbox.mutex_); - client.total_capacity_bytes = record->total_capacity_bytes; - client.enable_offloading = record->mailbox.enable_offloading_; - client.pending_offloads.insert( - record->mailbox.pending_offloads_.begin(), - record->mailbox.pending_offloads_.end()); - } - state.emplace(client_access.client_id, std::move(client)); - } - return state; -} - -void LocalSsdManager::RestorePersistedState(LocalSsdPersistedState state) { - ClientMap clients; - for (auto& [client_id, persisted] : state) { - auto record = - std::make_shared(persisted.enable_offloading); - record->total_capacity_bytes = persisted.total_capacity_bytes; - record->mailbox.pending_offloads_.insert( - std::make_move_iterator(persisted.pending_offloads.begin()), - std::make_move_iterator(persisted.pending_offloads.end())); - clients.emplace(client_id, std::move(record)); - } - ClientMap previous; - { - std::unique_lock lock(mutex_); - previous.swap(clients_); - clients_ = std::move(clients); - } - WaitForOperations(previous); -} - void LocalSsdManager::Clear() { ClientMap clients; { diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 488c018e3f..914c470208 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -46,7 +46,6 @@ #include "ha/oplog/oplog_batch_storage.h" #include "ha/oplog/ordered_oplog_writer.h" #include "types.h" -#include "serialize/serializer.h" #include "common/zstd_util.h" #include "common/file_util.h" #include "storage/distributed/shard_allocator.h" diff --git a/mooncake-store/src/offset_allocator.cpp b/mooncake-store/src/offset_allocator.cpp index a3cde866aa..a8e2ff1dee 100644 --- a/mooncake-store/src/offset_allocator.cpp +++ b/mooncake-store/src/offset_allocator.cpp @@ -544,279 +544,6 @@ static uint64_t calculateMultiplier(size_t size) { return multiplier_bits; } -tl::expected OffsetAllocatorSnapshot::ValidateNodeCapacity( - uint32_t size, uint32_t current_capacity, uint32_t max_capacity, - uint32_t free_offset) { - if (size == 0 || current_capacity == 0 || current_capacity > max_capacity || - max_capacity > kMaxNodes || free_offset > current_capacity) { - return tl::make_unexpected(fmt::format( - "invalid layout capacity fields: size={}, current_capacity={}, " - "max_capacity={}, free_offset={} (node limit={})", - size, current_capacity, max_capacity, free_offset, kMaxNodes)); - } - return {}; -} - -tl::expected OffsetAllocatorSnapshot::ValidateHeader() - const { - if (!layout) { - return tl::make_unexpected(std::string("layout is missing")); - } - if (capacity == 0 || capacity > std::numeric_limits::max() || - base > std::numeric_limits::max() - capacity) { - return tl::make_unexpected(fmt::format( - "invalid allocator bounds: base={}, capacity={}", base, capacity)); - } - // Every shift below, including usage accounting, is gated by this check. - const uint64_t expected_multiplier = calculateMultiplier(capacity); - if (multiplier_bits >= 64 || multiplier_bits != expected_multiplier) { - return tl::make_unexpected(fmt::format( - "invalid multiplier_bits: got {}, expected {} for capacity {}", - multiplier_bits, expected_multiplier, capacity)); - } - if (layout->m_size != (capacity >> multiplier_bits)) { - return tl::make_unexpected(fmt::format( - "layout size {} does not match capacity {} with multiplier_bits {}", - layout->m_size, capacity, multiplier_bits)); - } - if (auto valid = - ValidateNodeCapacity(layout->m_size, layout->m_current_capacity, - layout->m_max_capacity, layout->m_freeOffset); - !valid) { - return valid; - } - if (layout->m_nodes.size() != layout->m_current_capacity || - layout->m_freeNodes.size() != layout->m_current_capacity) { - return tl::make_unexpected(fmt::format( - "layout storage does not match current_capacity {}: nodes={}, " - "free_nodes={}", - layout->m_current_capacity, layout->m_nodes.size(), - layout->m_freeNodes.size())); - } - return {}; -} - -tl::expected -OffsetAllocatorSnapshot::ValidateUsedNodes( - std::vector& state) const { - LayoutUsage usage; - for (uint32_t i = 0; i < layout->m_current_capacity; ++i) { - const auto& node = layout->m_nodes[i]; - if (!node.used) continue; - if (node.dataSize == 0 || node.dataSize > layout->m_size) { - return tl::make_unexpected(fmt::format( - "used node {} has invalid size {}", i, node.dataSize)); - } - uint32_t min_units = node.dataSize; -#ifndef OFFSET_ALLOCATOR_NOT_ROUND_UP - const uint32_t bin = SmallFloat::uintToFloatRoundDown(node.dataSize); - if (SmallFloat::floatToUint(bin) != node.dataSize) { - return tl::make_unexpected( - fmt::format("used node {} size {} is not a bin boundary", i, - node.dataSize)); - } - // allocate() rounds first to a unit and then up to a bin. A request - // landing in this bin must be greater than the preceding bin boundary. - min_units = SmallFloat::floatToUint(bin - 1) + 1; -#endif - const uint64_t max_bytes = static_cast(node.dataSize) - << multiplier_bits; - if (max_bytes > capacity - usage.max_used_bytes) { - return tl::make_unexpected(fmt::format( - "used node {} makes occupied bytes exceed capacity {}", i, - capacity)); - } - usage.max_used_bytes += max_bytes; - usage.min_used_bytes += - (static_cast(min_units - 1) << multiplier_bits) + 1; - ++usage.used_nodes; - state[i] = NodeState::kUsed; - } - return usage; -} - -tl::expected OffsetAllocatorSnapshot::ValidateBins( - std::vector& state) const { - constexpr uint32_t kUnused = __Allocator::Node::unused; - uint8_t expected_used_bins[NUM_TOP_BINS] = {}; - uint32_t expected_used_bins_top = 0; - uint32_t bin_nodes = 0; - uint32_t free_storage = 0; - // Classification rejects cycles, duplicate membership and used nodes in - // free bins; the previous link and size class must agree with each list. - for (uint32_t bin = 0; bin < NUM_LEAF_BINS; ++bin) { - uint32_t node_index = layout->m_binIndices[bin]; - if (node_index == kUnused) continue; - const uint32_t top_bin = bin >> TOP_BINS_INDEX_SHIFT; - expected_used_bins[top_bin] |= uint8_t{1} - << (bin & LEAF_BINS_INDEX_MASK); - expected_used_bins_top |= uint32_t{1} << top_bin; - uint32_t previous = kUnused; - while (node_index != kUnused) { - if (node_index >= state.size() || - state[node_index] != NodeState::kUnclassified) { - return tl::make_unexpected(fmt::format( - "bin {} references out-of-range or already classified node " - "{}", - bin, node_index)); - } - const auto& node = layout->m_nodes[node_index]; - if (SmallFloat::uintToFloatRoundDown(node.dataSize) != bin || - node.binListPrev != previous) { - return tl::make_unexpected(fmt::format( - "bin {} node {} has inconsistent size class or previous " - "link", - bin, node_index)); - } - if (free_storage > - std::numeric_limits::max() - node.dataSize) { - return tl::make_unexpected(fmt::format( - "bin {} node {} overflows free storage", bin, node_index)); - } - state[node_index] = NodeState::kFree; - ++bin_nodes; - free_storage += node.dataSize; - previous = node_index; - node_index = node.binListNext; - } - } - if (layout->m_freeStorage != free_storage) { - return tl::make_unexpected( - fmt::format("free storage mismatch: got {}, expected {}", - layout->m_freeStorage, free_storage)); - } - if (layout->m_usedBinsTop != expected_used_bins_top) { - return tl::make_unexpected(std::string("top bin bitmap mismatch")); - } - for (uint32_t top = 0; top < NUM_TOP_BINS; ++top) { - if (layout->m_usedBins[top] != expected_used_bins[top]) { - return tl::make_unexpected( - fmt::format("leaf bin bitmap mismatch at top bin {}", top)); - } - } - return bin_nodes; -} - -tl::expected OffsetAllocatorSnapshot::ValidateFreeStack( - std::vector& state, uint32_t live_nodes) const { - for (uint32_t i = layout->m_freeOffset; i < state.size(); ++i) { - const uint32_t node_index = layout->m_freeNodes[i]; - if (node_index >= state.size() || - state[node_index] != NodeState::kUnclassified) { - return tl::make_unexpected(fmt::format( - "free stack slot {} references out-of-range or already " - "classified node {}", - i, node_index)); - } - state[node_index] = NodeState::kSpare; - } - if (layout->m_freeOffset != live_nodes) { - return tl::make_unexpected( - fmt::format("free stack offset mismatch: got {}, expected {}", - layout->m_freeOffset, live_nodes)); - } - for (uint32_t i = 0; i < state.size(); ++i) { - if (state[i] == NodeState::kUnclassified) { - return tl::make_unexpected(fmt::format( - "node {} is neither used, in a free bin, nor spare", i)); - } - } - return {}; -} - -tl::expected OffsetAllocatorSnapshot::ValidateNeighborChain( - const std::vector& state, uint32_t live_nodes) const { - constexpr uint32_t kUnused = __Allocator::Node::unused; - const auto is_live = [&state](uint32_t index) { - return index < state.size() && (state[index] == NodeState::kUsed || - state[index] == NodeState::kFree); - }; - uint32_t head = kUnused; - for (uint32_t i = 0; i < state.size(); ++i) { - if (!is_live(i) || layout->m_nodes[i].neighborPrev != kUnused) continue; - if (head != kUnused) { - return tl::make_unexpected( - std::string("neighbor chain has multiple heads")); - } - head = i; - } - if (head == kUnused) { - return tl::make_unexpected(std::string("neighbor chain has no head")); - } - std::vector walked(state.size(), 0); - uint32_t chain_nodes = 0; - uint32_t offset = 0; - for (uint32_t i = head; i != kUnused; i = layout->m_nodes[i].neighborNext) { - if (!is_live(i) || walked[i]) { - return tl::make_unexpected(fmt::format( - "neighbor chain references invalid or repeated node {}", i)); - } - walked[i] = 1; - const auto& node = layout->m_nodes[i]; - if (node.dataOffset != offset || node.dataSize == 0 || - node.dataSize > layout->m_size - offset) { - return tl::make_unexpected(fmt::format( - "neighbor node {} has invalid geometry: offset={}, size={}, " - "expected_offset={}", - i, node.dataOffset, node.dataSize, offset)); - } - offset += node.dataSize; - ++chain_nodes; - if (node.neighborNext != kUnused && - (!is_live(node.neighborNext) || - layout->m_nodes[node.neighborNext].neighborPrev != i)) { - return tl::make_unexpected(fmt::format( - "neighbor node {} has an inconsistent next link", i)); - } - if (node.neighborPrev != kUnused && - (!is_live(node.neighborPrev) || - layout->m_nodes[node.neighborPrev].neighborNext != i)) { - return tl::make_unexpected(fmt::format( - "neighbor node {} has an inconsistent previous link", i)); - } - } - if (offset != layout->m_size || chain_nodes != live_nodes) { - return tl::make_unexpected( - fmt::format("neighbor chain does not cover the layout: " - "bytes={}/{}, nodes={}/{}", - offset, layout->m_size, chain_nodes, live_nodes)); - } - return {}; -} - -tl::expected OffsetAllocatorSnapshot::ValidateAccounting( - const LayoutUsage& usage) const { - if (allocated_num != usage.used_nodes) { - return tl::make_unexpected( - fmt::format("allocated_num mismatch: got {}, expected {}", - allocated_num, usage.used_nodes)); - } - if (allocated_size < usage.min_used_bytes || - allocated_size > usage.max_used_bytes) { - return tl::make_unexpected(fmt::format( - "allocated_size {} is outside layout's requested-byte range [{}, " - "{}]", - allocated_size, usage.min_used_bytes, usage.max_used_bytes)); - } - return {}; -} - -tl::expected OffsetAllocatorSnapshot::Validate() const { - if (auto valid = ValidateHeader(); !valid) return valid; - std::vector state(layout->m_current_capacity, - NodeState::kUnclassified); - auto usage = ValidateUsedNodes(state); - if (!usage) return tl::make_unexpected(usage.error()); - auto bin_nodes = ValidateBins(state); - if (!bin_nodes) return tl::make_unexpected(bin_nodes.error()); - const uint32_t live_nodes = - static_cast(usage->used_nodes) + *bin_nodes; - if (auto valid = ValidateFreeStack(state, live_nodes); !valid) return valid; - if (auto valid = ValidateNeighborChain(state, live_nodes); !valid) - return valid; - return ValidateAccounting(*usage); -} - // Thread-safe OffsetAllocator implementation std::shared_ptr OffsetAllocator::create(uint64_t base, size_t size, @@ -827,29 +554,6 @@ std::shared_ptr OffsetAllocator::create(uint64_t base, new OffsetAllocator(base, size, init_capacity, max_capacity)); } -OffsetAllocatorSnapshot OffsetAllocator::CaptureSnapshot() const { - // The snapshot owns its layout: it can outlive this allocator and must - // never observe later mutations, so the bin/node state is copied here. - // The copy constructor is private; OffsetAllocator is a friend. - std::unique_ptr<__Allocator> layout(new __Allocator(*m_allocator)); - return {m_base, m_multiplier_bits, m_capacity, - m_allocated_size, m_allocated_num, std::move(layout)}; -} -tl::expected, std::string> -OffsetAllocator::Restore(OffsetAllocatorSnapshot snapshot) { - if (auto valid = snapshot.Validate(); !valid) { - LOG(ERROR) << "OffsetAllocator::Restore rejected snapshot: " - << valid.error(); - return tl::make_unexpected(valid.error()); - } - auto allocator = std::shared_ptr(new OffsetAllocator( - snapshot.base, snapshot.capacity, snapshot.multiplier_bits, - std::move(snapshot.layout))); - allocator->m_allocated_size = snapshot.allocated_size; - allocator->m_allocated_num = snapshot.allocated_num; - return allocator; -} - OffsetAllocator::OffsetAllocator(uint64_t base, size_t size, uint32 init_capacity, uint32 max_capacity) : m_base(base), @@ -862,21 +566,6 @@ OffsetAllocator::OffsetAllocator(uint64_t base, size_t size, std::memory_order_relaxed); } -OffsetAllocator::OffsetAllocator(uint64_t base, size_t size, - uint64_t multiplier_bits, - std::unique_ptr<__Allocator> allocator) - : m_allocator(std::move(allocator)), - m_base(base), - m_multiplier_bits(multiplier_bits), - m_capacity(size) { - const uint64_t largest_free_region = - m_allocator->storageReport().largestFreeRegion << m_multiplier_bits; - m_largest_free_region.store(largest_free_region, std::memory_order_relaxed); - const uint64_t allocator_capacity = - static_cast(m_allocator->m_size) << m_multiplier_bits; - m_largest_free_region_tightened = largest_free_region < allocator_capacity; -} - std::optional OffsetAllocator::allocate(size_t size) { if (size == 0) { return std::nullopt; diff --git a/mooncake-store/src/segment.cpp b/mooncake-store/src/segment.cpp index 04a85e5b9e..1b095c99e5 100644 --- a/mooncake-store/src/segment.cpp +++ b/mooncake-store/src/segment.cpp @@ -1,8 +1,6 @@ #include "segment.h" -#include "ha/snapshot/local_ssd_codec.h" #include "master_metric_manager.h" -#include "common/zstd_util.h" #include namespace mooncake { @@ -711,458 +709,6 @@ ErrorCode ScopedSegmentAccess::GetUnreadySegments( return ErrorCode::OK; } -ErrorCode SegmentView::GetSegment( - const std::shared_ptr& allocator, - Segment& segment) const { - // Check input parameters - if (!allocator) { - return ErrorCode::INVALID_PARAMS; - } - - // Iterate through mounted_segments_ to find matching allocator - for (const auto& [segment_id, mounted_segment] : - segment_manager_->mounted_segments_) { - if (mounted_segment.buf_allocator == allocator) { - segment = mounted_segment.segment; - return ErrorCode::OK; - } - } - - // No matching segment found - return ErrorCode::SEGMENT_NOT_FOUND; -} - -ErrorCode SegmentView::GetMountedSegment(const UUID& segment_id, - MountedSegment& mountedSegment) const { - // Check input parameters - if (segment_id.first == 0 && segment_id.second == 0) { - return ErrorCode::INVALID_PARAMS; - } - - // Find the specified segment_id in mounted_segments_ - auto it = segment_manager_->mounted_segments_.find(segment_id); - if (it == segment_manager_->mounted_segments_.end()) { - return ErrorCode::SEGMENT_NOT_FOUND; - } - - // Return the found segment - mountedSegment = it->second; - return ErrorCode::OK; -} - -tl::expected, SerializationError> -SegmentSerializer::Serialize(const LocalSsdPersistedState& local_ssd_state) { - if (!segment_manager_) { - return tl::unexpected( - SerializationError(ErrorCode::SERIALIZE_FAIL, - "serialize SegmentManager dependency is null")); - } - - if (segment_manager_->memory_allocator_ != BufferAllocatorType::OFFSET) { - return tl::unexpected(SerializationError( - ErrorCode::SERIALIZE_UNSUPPORTED, - "serialize SegmentManager memory_allocator_ is not offset")); - } - - msgpack::sbuffer sbuf; - msgpack::packer packer(&sbuf); - - // Create map containing all data - packer.pack_map( - 5); // 5 fields: memory_allocator, mounted_segments, client_segments, - // allocator_names, local_disk_segments - - // Serialize memory_allocator_ - packer.pack("ma"); - packer.pack(static_cast(segment_manager_->memory_allocator_)); - - // Serialize allocator_manager_'s names_ order - packer.pack("an"); // allocator_names - const auto& names = segment_manager_->allocator_manager_.getNames(); - packer.pack_array(names.size()); - for (const auto& name : names) { - packer.pack(name); - } - - // Serialize mounted_segments_ - // Sort UUIDs first to ensure deterministic serialization results - packer.pack("ms"); - packer.pack_map(segment_manager_->mounted_segments_.size()); - - // Collect all segment UUIDs and sort - std::vector sorted_segment_uuids; - sorted_segment_uuids.reserve(segment_manager_->mounted_segments_.size()); - for (const auto& pair : segment_manager_->mounted_segments_) { - sorted_segment_uuids.push_back(pair.first); - } - std::sort(sorted_segment_uuids.begin(), sorted_segment_uuids.end()); - - for (const auto& segment_uuid : sorted_segment_uuids) { - const auto& mounted_segment = - segment_manager_->mounted_segments_.at(segment_uuid); - std::string uuid_str = UuidToString(segment_uuid); - packer.pack(uuid_str); - - // Serialize MountedSegment object - auto result = - Serializer::serialize(mounted_segment, packer); - if (!result) { - return tl::unexpected(result.error()); - } - } - - // Serialize client_segments_ - // Sort client UUIDs first to ensure deterministic serialization results - packer.pack("cs"); - packer.pack_map(segment_manager_->client_segments_.size()); - - // Collect all client UUIDs and sort - std::vector sorted_client_uuids; - sorted_client_uuids.reserve(segment_manager_->client_segments_.size()); - for (const auto& pair : segment_manager_->client_segments_) { - sorted_client_uuids.push_back(pair.first); - } - std::sort(sorted_client_uuids.begin(), sorted_client_uuids.end()); - - for (const auto& client_uuid : sorted_client_uuids) { - const auto& segment_ids = - segment_manager_->client_segments_.at(client_uuid); - std::string client_uuid_str = UuidToString(client_uuid); - packer.pack(client_uuid_str); - - // Serialize segment IDs array - packer.pack_array(segment_ids.size()); - - for (const auto& segment_id : segment_ids) { - std::string segment_uuid_str = UuidToString(segment_id); - packer.pack(segment_uuid_str); - } - } - - packer.pack("ld"); // local_disk_segments - auto local_ssd_result = ha::LocalSsdCodec::Encode(local_ssd_state, packer); - if (!local_ssd_result) { - return tl::unexpected(local_ssd_result.error()); - } - - // Compress entire data - std::vector compressed_data = zstd_compress( - reinterpret_cast(sbuf.data()), sbuf.size(), 3); - - // Return compressed data - return std::vector( - std::make_move_iterator(compressed_data.begin()), - std::make_move_iterator(compressed_data.end())); -} - -tl::expected -SegmentSerializer::Deserialize(const std::vector& data) { - // Decompress data - std::vector decompressed_data; - try { - decompressed_data = zstd_decompress( - reinterpret_cast(data.data()), data.size()); - } catch (const std::exception& e) { - return tl::make_unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager failed to " - "decompress MessagePack data: " + - std::string(e.what()))); - } - - // Parse MessagePack data - msgpack::object_handle oh; - try { - oh = msgpack::unpack( - reinterpret_cast(decompressed_data.data()), - decompressed_data.size()); - } catch (const std::exception& e) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager failed to unpack MessagePack data: " + - std::string(e.what()))); - } - - const msgpack::object& obj = oh.get(); - - // Check if it's a map - if (obj.type != msgpack::type::MAP) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager invalid MessagePack " - "format: expected map")); - } - - if (!segment_manager_) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager dependency is null")); - } - - // Convert MessagePack map to regular map, use pointers for values to avoid - // copying - std::map fields_map; - for (uint32_t i = 0; i < obj.via.map.size; ++i) { - const msgpack::object& key_obj = obj.via.map.ptr[i].key; - const msgpack::object& value_obj = obj.via.map.ptr[i].val; - - // Ensure key is a string - if (key_obj.type == msgpack::type::STR) { - std::string key(key_obj.via.str.ptr, key_obj.via.str.size); - fields_map.emplace(std::move(key), &value_obj); - } - } - - auto ld_it = fields_map.find("ld"); - auto local_ssd_state = ha::LocalSsdCodec::Decode( - ld_it == fields_map.end() ? nullptr : ld_it->second); - if (!local_ssd_state) { - return tl::unexpected(local_ssd_state.error()); - } - - // Do not mutate SegmentManager until the independent LocalSSD payload has - // been fully validated. - segment_manager_->mounted_segments_.clear(); - segment_manager_->client_segments_.clear(); - segment_manager_->segments_by_host_.clear(); - - // Process fields in order - // 1. First process memory_allocator_ - auto allocator_it = fields_map.find("ma"); - if (allocator_it != fields_map.end()) { - const msgpack::object* value_obj = allocator_it->second; - if (value_obj->type != msgpack::type::POSITIVE_INTEGER) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager memory_allocator is not int")); - } - - auto allocatorType = - static_cast(value_obj->as()); - - // Note: Type must match, only OffsetAllocator is supported - if (allocatorType != segment_manager_->memory_allocator_) { - LOG(ERROR) << "deserialize memory allocator type doesn't match " - "current setting"; - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager memory allocator type doesn't " - "match current setting")); - } - } - - // 1.5 Parse allocator_names (saved original names_ order) - std::vector saved_allocator_names; - auto allocator_names_it = fields_map.find("an"); - if (allocator_names_it != fields_map.end()) { - const msgpack::object* value_obj = allocator_names_it->second; - if (value_obj->type != msgpack::type::ARRAY) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager allocator_names is not array")); - } - - saved_allocator_names.reserve(value_obj->via.array.size); - for (uint32_t i = 0; i < value_obj->via.array.size; ++i) { - const msgpack::object& name_obj = value_obj->via.array.ptr[i]; - if (name_obj.type != msgpack::type::STR) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager allocator_name is not string")); - } - saved_allocator_names.emplace_back(name_obj.via.str.ptr, - name_obj.via.str.size); - } - } - - // 2. Process mounted_segments_ - auto mounted_segments_it = fields_map.find("ms"); - if (mounted_segments_it != fields_map.end()) { - const msgpack::object* value_obj = mounted_segments_it->second; - if (value_obj->type != msgpack::type::MAP) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager mounted_segments is not map")); - } - - // Iterate through all mounted segments - for (uint32_t j = 0; j < value_obj->via.map.size; ++j) { - const msgpack::object& segment_key = value_obj->via.map.ptr[j].key; - const msgpack::object& segment_value = - value_obj->via.map.ptr[j].val; - - // Ensure key is a string (UUID) - if (segment_key.type != msgpack::type::STR) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager " - "mounted_segments key is not string")); - } - - std::string uuid_str(segment_key.via.str.ptr, - segment_key.via.str.size); - UUID segment_uuid; - if (!StringToUuid(uuid_str, segment_uuid)) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize SegmentManager mounted_segments " - "uuid {} is invalid", - uuid_str))); - } - - // Deserialize MountedSegment object - auto result = - Serializer::deserialize(segment_value); - if (!result) { - return tl::unexpected(result.error()); - } - - segment_manager_->mounted_segments_.emplace( - segment_uuid, std::move(result.value())); - } - } - - // 3. Process client_segments_ - auto client_segments_it = fields_map.find("cs"); - if (client_segments_it != fields_map.end()) { - const msgpack::object* value_obj = client_segments_it->second; - if (value_obj->type != msgpack::type::MAP) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager client_segments is not map")); - } - - // Iterate through all client segments - for (uint32_t j = 0; j < value_obj->via.map.size; ++j) { - const msgpack::object& client_key = value_obj->via.map.ptr[j].key; - const msgpack::object& client_value = value_obj->via.map.ptr[j].val; - - // Ensure key is a string (UUID) - if (client_key.type != msgpack::type::STR) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager " - "client_segments key is not string")); - } - - std::string client_uuid_str(client_key.via.str.ptr, - client_key.via.str.size); - UUID client_uuid; - if (!StringToUuid(client_uuid_str, client_uuid)) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize SegmentManager client_segments " - "client uuid {} is invalid", - client_uuid_str))); - } - - // Ensure value is an array - if (client_value.type != msgpack::type::ARRAY) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager " - "client_segments value is not array")); - } - - // Deserialize segment IDs array - std::vector segment_ids; - segment_ids.reserve(client_value.via.array.size); - - for (uint32_t k = 0; k < client_value.via.array.size; ++k) { - const msgpack::object& segment_id_obj = - client_value.via.array.ptr[k]; - - if (segment_id_obj.type != msgpack::type::STR) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize SegmentManager segment_id is not string")); - } - - std::string segment_uuid_str(segment_id_obj.via.str.ptr, - segment_id_obj.via.str.size); - UUID segment_uuid; - if (!StringToUuid(segment_uuid_str, segment_uuid)) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize SegmentManager segment_id {} " - "is invalid", - segment_uuid_str))); - } - - segment_ids.push_back(segment_uuid); - } - - segment_manager_->client_segments_.emplace(client_uuid, - std::move(segment_ids)); - } - } - - // Restore allocator_manager_ based on mounted_segments_ and saved names - // order - segment_manager_->allocator_manager_ = AllocatorManager(); - for (auto& [segment_id, mounted_segment] : - segment_manager_->mounted_segments_) { - (void)segment_id; - if (!mounted_segment.buf_allocator) { - continue; - } - mounted_segment.allocator_registration = - std::shared_ptr( - new SegmentAllocatorRegistration(mounted_segment.buf_allocator, - nullptr)); - if (mounted_segment.status != SegmentStatus::OK) { - mounted_segment.allocator_registration->SetAllocatable(false); - } - } - - // Add allocators in saved original order - for (const auto& name : saved_allocator_names) { - for (auto& pair : segment_manager_->mounted_segments_) { - MountedSegment& mounted_segment = pair.second; - if (mounted_segment.segment.name == name && - mounted_segment.status == SegmentStatus::OK && - mounted_segment.buf_allocator) { - segment_manager_->allocator_manager_.addRegistration( - name, mounted_segment.allocator_registration); - } - } - } - - // Rebuild segment indexes based on client_segments_ and mounted_segments_ - segment_manager_->client_by_name_.clear(); - segment_manager_->segment_id_by_name_.clear(); - segment_manager_->segments_by_host_.clear(); - for (const auto& [client_id, segment_ids] : - segment_manager_->client_segments_) { - for (const auto& segment_id : segment_ids) { - auto it = segment_manager_->mounted_segments_.find(segment_id); - if (it != segment_manager_->mounted_segments_.end()) { - segment_manager_->client_by_name_[it->second.segment.name] = - client_id; - segment_manager_->segment_id_by_name_[it->second.segment.name] = - segment_id; - if (it->second.status == SegmentStatus::OK && - it->second.buf_allocator) { - AddHostSegment(segment_manager_->segments_by_host_, - it->second.segment); - } - } - } - } - - segment_manager_->AttachMountedUsageTrackers(); - return std::move(*local_ssd_state); -} - -void SegmentSerializer::Reset() { - segment_manager_->mounted_segments_.clear(); - segment_manager_->client_segments_.clear(); - segment_manager_->client_by_name_.clear(); - segment_manager_->segment_id_by_name_.clear(); - segment_manager_->segments_by_host_.clear(); - segment_manager_->allocator_manager_ = AllocatorManager(); -} - ErrorCode ScopedSegmentAccess::GetClientIdBySegmentName( const std::string& segment_name, UUID& client_id) const { auto it = segment_manager_->client_by_name_.find(segment_name); diff --git a/mooncake-store/src/segment/pool_mount.cpp b/mooncake-store/src/segment/pool_mount.cpp index a885122bae..94d945f43b 100644 --- a/mooncake-store/src/segment/pool_mount.cpp +++ b/mooncake-store/src/segment/pool_mount.cpp @@ -101,20 +101,18 @@ SegmentPool::WriteAccess::PrepareWithLiveAllocations( tl::expected SegmentPool::WriteAccess::PrepareAdopt( MountedRegion mounted, std::shared_ptr allocator, - bool account_capacity_metrics, AdoptMode mode) { + bool account_capacity_metrics) { mounted.kind = ClassifyRegion(mounted.segment); RegionDriver* driver = segment_pool_.GetDriver(mounted.kind); if (!driver) { return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); } - if (mode == AdoptMode::Insert) { - if (catalog_.Find(mounted.segment.id)) { - return tl::make_unexpected(ErrorCode::SEGMENT_ALREADY_EXISTS); - } - auto owner = catalog_.FindOwnerClientId(mounted.segment.name); - if (owner && *owner != mounted.client_id) { - return tl::make_unexpected(ErrorCode::INVALID_PARAMS); - } + if (catalog_.Find(mounted.segment.id)) { + return tl::make_unexpected(ErrorCode::SEGMENT_ALREADY_EXISTS); + } + auto owner = catalog_.FindOwnerClientId(mounted.segment.name); + if (owner && *owner != mounted.client_id) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } auto resource = driver->PrepareAdopt(MakeResourceSpec(mounted.segment), std::move(allocator)); diff --git a/mooncake-store/src/segment/snapshot.cpp b/mooncake-store/src/segment/snapshot.cpp deleted file mode 100644 index 815a123ead..0000000000 --- a/mooncake-store/src/segment/snapshot.cpp +++ /dev/null @@ -1,123 +0,0 @@ -#include "segment/snapshot.h" - -#include -#include -#include - -#include "segment/pool.h" -#include "segment/pool_write_access.h" - -namespace mooncake { - -tl::expected SegmentPool::CaptureSnapshot() - const { - const auto* driver = GetDriver(RegionKind::HOST_MEMORY); - if (!driver || driver->allocator_type() != BufferAllocatorType::OFFSET) { - return tl::make_unexpected(ErrorCode::SERIALIZE_UNSUPPORTED); - } - // Validate supported resources before copying any allocator layouts. - for (const auto& mounted : catalog_.Regions()) { - if (mounted.kind != RegionKind::HOST_MEMORY) { - return tl::make_unexpected(ErrorCode::SERIALIZE_UNSUPPORTED); - } - } - - SegmentPoolSnapshot snapshot; - placement_index_.GetActiveSegmentNames(AllocationCandidateKind::NATIVE, - snapshot.active_names); - snapshot.regions.reserve(catalog_.Regions().size()); - for (const auto& mounted : catalog_.Regions()) { - const auto* resource = GetResource(mounted); - if (!resource || !resource->allocator()) { - return tl::make_unexpected(ErrorCode::SERIALIZE_FAIL); - } - const auto* allocator = dynamic_cast( - resource->allocator().get()); - if (!allocator) { - return tl::make_unexpected(ErrorCode::SERIALIZE_UNSUPPORTED); - } - snapshot.regions.push_back({mounted, allocator->CaptureSnapshot()}); - } - std::sort(snapshot.regions.begin(), snapshot.regions.end(), - [](const auto& lhs, const auto& rhs) { - return lhs.mounted.segment.id < rhs.mounted.segment.id; - }); - return snapshot; -} - -tl::expected SegmentPool::RestoreSnapshot( - SegmentPoolSnapshot snapshot, bool account_capacity_metrics) { - auto access = AcquireWriteAccess(); - const auto* driver = GetDriver(RegionKind::HOST_MEMORY); - if (!driver || - snapshot.memory_allocator_type != BufferAllocatorType::OFFSET || - driver->allocator_type() != snapshot.memory_allocator_type) { - return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); - } - std::set ids; - std::map owners; - std::set active_names; - for (const auto& region : snapshot.regions) { - const auto& mounted = region.mounted; - if (!IsKnownSegmentStatus(mounted.status)) { - return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); - } - auto [owner, inserted] = - owners.emplace(mounted.segment.name, mounted.client_id); - if (!ids.insert(mounted.segment.id).second || - (!inserted && owner->second != mounted.client_id) || - mounted.kind != RegionKind::HOST_MEMORY || - region.allocator.segment_name != mounted.segment.name || - region.allocator.base != mounted.segment.base || - region.allocator.capacity != mounted.segment.size || - region.allocator.transport_endpoint != - mounted.segment.te_endpoint) { - return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); - } - if (mounted.status == SegmentStatus::OK) - active_names.insert(mounted.segment.name); - } - if (active_names != std::set(snapshot.active_names.begin(), - snapshot.active_names.end()) || - active_names.size() != snapshot.active_names.size()) { - return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); - } - std::vector prepared; - prepared.reserve(snapshot.regions.size()); - // Preserve the active-name order even for snapshots captured in ID order. - auto prepare = - [&](RegionSnapshot& region) -> tl::expected { - auto allocator = - OffsetBufferAllocator::Restore(std::move(region.allocator)); - if (!allocator) return tl::make_unexpected(allocator.error()); - auto txn = access.PrepareAdopt(region.mounted, std::move(*allocator), - account_capacity_metrics, - WriteAccess::AdoptMode::ReplacePool); - if (!txn) return tl::make_unexpected(txn.error()); - prepared.push_back(std::move(*txn)); - return {}; - }; - for (const auto& name : snapshot.active_names) { - for (auto& region : snapshot.regions) { - if (region.mounted.status == SegmentStatus::OK && - region.mounted.segment.name == name) { - auto result = prepare(region); - if (!result) return result; - } - } - } - for (auto& region : snapshot.regions) { - if (region.mounted.status != SegmentStatus::OK) { - auto result = prepare(region); - if (!result) return result; - } - } - access.Clear(); - for (auto& txn : prepared) { - const auto result = txn.Commit(access); - if (result != ErrorCode::OK) return tl::make_unexpected(result); - } - return {}; -} - -} // namespace mooncake diff --git a/mooncake-store/src/serialize/serializer.cpp b/mooncake-store/src/serialize/serializer.cpp deleted file mode 100644 index 2006d22e4c..0000000000 --- a/mooncake-store/src/serialize/serializer.cpp +++ /dev/null @@ -1,1080 +0,0 @@ -#include -#include -#include -#include - -#include "serialize/serializer.h" -#include "ha/snapshot/allocator_snapshot_codec.h" -#include "offset_allocator/offset_allocator.h" -#include "types.h" -#include "master_service.h" -#include "common/zstd_util.h" - -namespace mooncake { - -// Node serialization size constants for offset_allocator::__Allocator -constexpr size_t OFFSET_ALLOCATOR_NODE_BOOL_SIZE = 1; -constexpr size_t OFFSET_ALLOCATOR_NODE_UINT32_COUNT = 6; -constexpr size_t OFFSET_ALLOCATOR_NODE_SERIALIZED_SIZE = - OFFSET_ALLOCATOR_NODE_BOOL_SIZE + - OFFSET_ALLOCATOR_NODE_UINT32_COUNT * sizeof(uint32_t); - -// __Allocator serialize_msgpack -tl::expected -Serializer::serialize( - const offset_allocator::__Allocator &allocator, MsgpackPacker &packer) { - // Use array instead of map for more compact storage - // Array order is consistent with deserialize_msgpack - packer.pack_array(10); - - // Basic properties (packed in order) - packer.pack(allocator.m_size); - packer.pack(allocator.m_current_capacity); - packer.pack(allocator.m_max_capacity); - packer.pack(allocator.m_freeStorage); - packer.pack(allocator.m_usedBinsTop); - - // usedBins array - packer.pack_array(offset_allocator::NUM_TOP_BINS); - for (unsigned char m_usedBin : allocator.m_usedBins) { - packer.pack(m_usedBin); - } - - // binIndex array - packer.pack_array(offset_allocator::NUM_LEAF_BINS); - for (unsigned int m_binIndex : allocator.m_binIndices) { - packer.pack(m_binIndex); - } - - // nodes data serialization and compression - // Each node serializes to: 1 byte (used flag) + 6 * sizeof(uint32_t) - // (fields) - std::vector serialized_nodes; - serialized_nodes.reserve(allocator.m_max_capacity * - OFFSET_ALLOCATOR_NODE_SERIALIZED_SIZE); - - for (uint32_t i = 0; i < allocator.m_current_capacity; i++) { - const auto &node = allocator.m_nodes[i]; - serialized_nodes.push_back(node.used ? 1 : 0); - SerializationHelper::serializeUint32(node.dataOffset, serialized_nodes); - SerializationHelper::serializeUint32(node.dataSize, serialized_nodes); - SerializationHelper::serializeUint32(node.binListPrev, - serialized_nodes); - SerializationHelper::serializeUint32(node.binListNext, - serialized_nodes); - SerializationHelper::serializeUint32(node.neighborPrev, - serialized_nodes); - SerializationHelper::serializeUint32(node.neighborNext, - serialized_nodes); - } - - try { - std::vector compressed_nodes = - zstd_compress(serialized_nodes, 3); - packer.pack(compressed_nodes); - - // freeNodes data serialization and compression - std::vector serialized_free_nodes; - serialized_free_nodes.reserve(allocator.m_current_capacity * 4); - - for (uint32_t i = 0; i < allocator.m_current_capacity; i++) { - SerializationHelper::serializeUint32(allocator.m_freeNodes[i], - serialized_free_nodes); - } - - std::vector compressed_free_nodes = - zstd_compress(serialized_free_nodes, 3); - packer.pack(compressed_free_nodes); - } catch (const std::exception &e) { - return tl::unexpected(SerializationError( - ErrorCode::SERIALIZE_FAIL, - std::string( - "offset_allocator::__Allocator, error compressing nodes: ") + - e.what())); - } - - // freeOffset - packer.pack(allocator.m_freeOffset); - - return {}; -} - -// __Allocator deserialize_msgpack -tl::expected, SerializationError> -Serializer::deserialize( - const msgpack::object &obj) { - // Check if object type is array (consistent with serialize_msgpack) - if (obj.type != msgpack::type::ARRAY) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize offset_allocator::__Allocator " - "invalid msgpack data, expected array")); - } - - // Verify array size is correct - if (obj.via.array.size != 10) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::__Allocator invalid " - "array size: expected 10, got {}", - obj.via.array.size))); - } - - auto *array_items = obj.via.array.ptr; - size_t index = 0; - - // Parse every scalar before materializing or decompressing anything: the - // allocator constructor reserves max_capacity nodes, so corrupt capacities - // and a corrupt freeOffset must be rejected against the node bound before a - // payload can request large storage. - uint32_t size = 0; - uint32_t current_capacity = 0; - uint32_t max_capacity = 0; - uint32_t free_storage = 0; - uint32_t used_bins_top = 0; - uint32_t free_offset = 0; - try { - size = array_items[index++].as(); - current_capacity = array_items[index++].as(); - max_capacity = array_items[index++].as(); - free_storage = array_items[index++].as(); - used_bins_top = array_items[index++].as(); - // freeOffset is packed last; reading it here keeps the whole scalar - // header validated before any node payload is touched. - free_offset = array_items[9].as(); - } catch (const std::exception &e) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::__Allocator invalid " - "scalar field: {}", - e.what()))); - } - - if (auto valid = - offset_allocator::OffsetAllocatorSnapshot::ValidateNodeCapacity( - size, current_capacity, max_capacity, free_offset); - !valid) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::__Allocator {}", - valid.error()))); - } - - std::array used_bins{}; - std::array - bin_indices{}; - try { - // Deserialize usedBins array - const auto &used_bins_array = array_items[index++]; - if (used_bins_array.type != msgpack::type::ARRAY) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize offset_allocator::__Allocator " - "usedBins is not an array")); - } - - if (used_bins_array.via.array.size != offset_allocator::NUM_TOP_BINS) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format( - "deserialize offset_allocator::__Allocator usedBins " - "invalid size: expected {}, got {}", - offset_allocator::NUM_TOP_BINS, - used_bins_array.via.array.size))); - } - - for (uint32_t i = 0; i < offset_allocator::NUM_TOP_BINS; i++) { - used_bins[i] = used_bins_array.via.array.ptr[i].as(); - } - - // Deserialize binIndices array - const auto &bin_indices_array = array_items[index++]; - if (bin_indices_array.type != msgpack::type::ARRAY) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize offset_allocator::__Allocator " - "binIndices is not an array")); - } - - if (bin_indices_array.via.array.size != - offset_allocator::NUM_LEAF_BINS) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format( - "deserialize offset_allocator::__Allocator binIndices " - "invalid size: expected {}, got {}", - offset_allocator::NUM_LEAF_BINS, - bin_indices_array.via.array.size))); - } - - for (uint32_t i = 0; i < offset_allocator::NUM_LEAF_BINS; i++) { - const uint32_t bin_index = - bin_indices_array.via.array.ptr[i].as(); - if (bin_index != offset_allocator::__Allocator::Node::unused && - bin_index >= current_capacity) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::__Allocator bin " - "index {} is out of range for capacity {}", - bin_index, current_capacity))); - } - bin_indices[i] = bin_index; - } - } catch (const std::exception &e) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::__Allocator invalid " - "bitmap field: {}", - e.what()))); - } - - std::vector serialized_nodes; - try { - // Deserialize compressed nodes data - const auto &nodes_bin = array_items[index++]; - if (nodes_bin.type != msgpack::type::BIN) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize offset_allocator::__Allocator " - "nodes data is not binary")); - } - - // The payload must describe exactly the nodes it will populate, so the - // decompression is bounded by that size: a frame header can then never - // request a far larger buffer. - const size_t expected_size = static_cast(current_capacity) * - OFFSET_ALLOCATOR_NODE_SERIALIZED_SIZE; - serialized_nodes = zstd_decompress( - reinterpret_cast(nodes_bin.via.bin.ptr), - nodes_bin.via.bin.size, expected_size); - if (serialized_nodes.size() != expected_size) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::__Allocator invalid " - "serialized nodes data size: expected {}, actual " - "size: {}", - expected_size, serialized_nodes.size()))); - } - } catch (const std::exception &e) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::__Allocator error " - "decompressing nodes: {}", - e.what()))); - } - - std::vector free_nodes; - try { - // Deserialize compressed freeNodes data - const auto &free_nodes_bin = array_items[index++]; - if (free_nodes_bin.type != msgpack::type::BIN) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize offset_allocator::__Allocator " - "freeNodes data is not binary")); - } - - // Bounded like the node payload, then required to be exact. - const size_t expected_size = static_cast(current_capacity) * 4; - const std::vector serialized_free_nodes = zstd_decompress( - reinterpret_cast(free_nodes_bin.via.bin.ptr), - free_nodes_bin.via.bin.size, expected_size); - if (serialized_free_nodes.size() != expected_size) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format( - "deserialize offset_allocator::__Allocator invalid " - "serialized free nodes data size: expected {}, actual {}", - expected_size, serialized_free_nodes.size()))); - } - - // Deserialize freeNodes array in standardized format - free_nodes.resize(current_capacity); - for (uint32_t i = 0; i < current_capacity; i++) { - free_nodes[i] = SerializationHelper::deserializeUint32( - &serialized_free_nodes[static_cast(i) * 4]); - } - } catch (const std::exception &e) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::__Allocator error " - "processing free nodes: {}", - e.what()))); - } - - // Capacities and payload sizes are consistent: both the reservation and - // populated node storage are bounded, and every copy stays in range. - try { - auto allocator = std::make_unique( - size, current_capacity, max_capacity); - allocator->m_freeStorage = free_storage; - allocator->m_usedBinsTop = used_bins_top; - allocator->m_freeOffset = free_offset; - for (uint32_t i = 0; i < offset_allocator::NUM_TOP_BINS; i++) { - allocator->m_usedBins[i] = used_bins[i]; - } - for (uint32_t i = 0; i < offset_allocator::NUM_LEAF_BINS; i++) { - allocator->m_binIndices[i] = bin_indices[i]; - } - - // Deserialize nodes array in standardized format - size_t offset = 0; - for (uint32_t i = 0; i < current_capacity; i++) { - // Deserialize bool field - allocator->m_nodes[i].used = (serialized_nodes[offset++] != 0); - - // Deserialize uint32_t fields - allocator->m_nodes[i].dataOffset = - SerializationHelper::deserializeUint32( - &serialized_nodes[offset]); - offset += 4; - allocator->m_nodes[i].dataSize = - SerializationHelper::deserializeUint32( - &serialized_nodes[offset]); - offset += 4; - allocator->m_nodes[i].binListPrev = - SerializationHelper::deserializeUint32( - &serialized_nodes[offset]); - offset += 4; - allocator->m_nodes[i].binListNext = - SerializationHelper::deserializeUint32( - &serialized_nodes[offset]); - offset += 4; - allocator->m_nodes[i].neighborPrev = - SerializationHelper::deserializeUint32( - &serialized_nodes[offset]); - offset += 4; - allocator->m_nodes[i].neighborNext = - SerializationHelper::deserializeUint32( - &serialized_nodes[offset]); - offset += 4; - } - for (uint32_t i = 0; i < current_capacity; i++) { - allocator->m_freeNodes[i] = free_nodes[i]; - } - - return allocator; - } catch (const std::exception &e) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::__Allocator error " - "constructing layout: {}", - e.what()))); - } -} - -// serialize_msgpack -tl::expected -Serializer::serialize( - const offset_allocator::OffsetAllocator &allocator, MsgpackPacker &packer) { - packer.pack_array(6); - - // Serialize basic members (packed in order) - packer.pack(allocator.m_base); - packer.pack(allocator.m_multiplier_bits); - packer.pack(allocator.m_capacity); - - packer.pack(allocator.m_allocated_size); - packer.pack(allocator.m_allocated_num); - - // Serialize __Allocator - auto allocator_result = - Serializer::serialize( - *allocator.m_allocator, packer); - if (!allocator_result) { - return tl::unexpected(allocator_result.error()); - } - - return {}; -} - -// deserialize_msgpack -auto Serializer::deserialize( - const msgpack::object &obj) - -> tl::expected { - // Check if object type is array (consistent with serialize_msgpack) - if (obj.type != msgpack::type::ARRAY) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize offset_allocator::OffsetAllocator " - "invalid msgpack data, expected array")); - } - - // Verify array size is correct (should have 6 elements, consistent with - // serialize_msgpack) - if (obj.via.array.size != 6) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::OffsetAllocator invalid " - "array size: expected 6, got {}", - obj.via.array.size))); - } - - auto *array_items = obj.via.array.ptr; - size_t index = 0; - - try { - // Deserialize basic properties - uint64_t base = array_items[index++].as(); - uint64_t multiplier_bits = array_items[index++].as(); - uint64_t capacity = array_items[index++].as(); - - uint64_t allocated_size = array_items[index++].as(); - uint64_t allocated_num = array_items[index++].as(); - - // Deserialize __Allocator - auto allocator_result = - Serializer::deserialize( - array_items[index++]); - if (!allocator_result) { - return tl::unexpected(allocator_result.error()); - } - - // Install through the same factory that snapshot restore uses: it runs - // the full layout validation exactly once and returns the detailed - // reason. - offset_allocator::OffsetAllocatorSnapshot snapshot{ - .base = base, - .multiplier_bits = multiplier_bits, - .capacity = capacity, - .allocated_size = allocated_size, - .allocated_num = allocated_num, - .layout = std::move(allocator_result.value()), - }; - auto offset_allocator = - offset_allocator::OffsetAllocator::Restore(std::move(snapshot)); - if (!offset_allocator) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format( - "deserialize offset_allocator::OffsetAllocator rejected " - "its state: {}", - offset_allocator.error()))); - } - - return *offset_allocator; - } catch (const std::exception &e) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format( - "deserialize offset_allocator::OffsetAllocator failed: {}", - e.what()))); - } -} - -tl::expected -Serializer::serialize( - const offset_allocator::OffsetAllocationHandle &handle, - MsgpackPacker &packer) { - packer.pack_array(3); - - // Serialize basic fields - packer.pack(handle.real_base); - packer.pack(handle.requested_size); - - // Serialize allocation struct fields - packer.pack_array(2); - packer.pack(handle.m_allocation.offset); - packer.pack(handle.m_allocation.metadata); - return {}; -} - -auto Serializer::deserialize( - const msgpack::object &obj, - const std::shared_ptr &allocator) - -> tl::expected { - // Check if object type is array (consistent with serialize_msgpack) - if (obj.type != msgpack::type::ARRAY) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize offset_allocator::OffsetAllocationHandle invalid " - "msgpack data, expected array")); - } - - // Verify array size is correct (should have 3 elements, consistent with - // serialize_msgpack) - if (obj.via.array.size != 3) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::OffsetAllocationHandle " - "invalid array size: expected 3, got {}", - obj.via.array.size))); - } - - auto *array_items = obj.via.array.ptr; - - // Deserialize basic properties - uint64_t real_base = array_items[0].as(); - uint64_t requested_size = array_items[1].as(); - - // Deserialize allocation struct - const auto &allocation_array = array_items[2]; - if (allocation_array.type != msgpack::type::ARRAY) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize offset_allocator::OffsetAllocationHandle allocation " - "is not an array")); - } - - if (allocation_array.via.array.size != 2) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize offset_allocator::OffsetAllocationHandle " - "allocation invalid size: expected 2, got {}", - allocation_array.via.array.size))); - } - - auto offset = allocation_array.via.array.ptr[0].as(); - auto metadata = allocation_array.via.array.ptr[1].as(); - offset_allocator::OffsetAllocation allocation(offset, metadata); - - // Create a new OffsetAllocationHandle object - auto handle = std::make_shared( - allocator, allocation, real_base, requested_size); - - return handle; -} - -tl::expected Serializer::serialize( - const AllocatedBuffer &buffer, const SegmentView &segment_view, - MsgpackPacker &packer) { - packer.pack_array(5); - - // Serialize basic properties - // packer.pack(buffer.segment_name_); - packer.pack(static_cast(buffer.size_)); - packer.pack(reinterpret_cast(buffer.buffer_ptr_)); - // packer.pack(static_cast(buffer.status)); - - if (buffer.allocator_.expired()) { - return tl::unexpected(SerializationError( - ErrorCode::SERIALIZE_FAIL, - fmt::format("buffer.allocator_.expired,buffer_ptr:{}", - buffer.buffer_ptr_))); - } - - // Get segment info - const auto &allocator = buffer.allocator_.lock(); - if (!allocator) { - return tl::unexpected(SerializationError( - ErrorCode::SERIALIZE_FAIL, - fmt::format("serialize AllocatedBuffer " - "buffer.allocator_.lock() fail,buffer_ptr:{}", - buffer.buffer_ptr_))); - } - - Segment segment; - ErrorCode ret = segment_view.GetSegment(allocator, segment); - if (ret != ErrorCode::OK) { - return tl::unexpected(SerializationError( - ErrorCode::SERIALIZE_FAIL, - fmt::format("serialize AllocatedBuffer " - "segment_view.GetSegment() fail ret={}", - static_cast(ret)))); - } - - packer.pack(UuidToString(segment.id)); - - // Serialize offset_handle_ (if exists) - if (buffer.offset_handle_.has_value()) { - // Mark offset_handle exists - packer.pack(true); - auto handle_result = - Serializer::serialize( - buffer.offset_handle_.value(), packer); - if (!handle_result) { - return tl::unexpected(handle_result.error()); - } - } else { - // Mark offset_handle does not exist - packer.pack(false); - packer.pack_nil(); - } - - return {}; -} - -auto Serializer::deserialize(const msgpack::object &obj, - const SegmentView &segment_view) - -> tl::expected { - // Check if object type is array (consistent with serialize_msgpack) - if (obj.type != msgpack::type::ARRAY) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize_msgpack AllocatedBuffer invalid " - "msgpack data, expected array")); - } - - // Verify array size is correct (should have 5 elements: size, buffer_ptr, - // segment_id, has_offset_handle, offset_handle) - if (obj.via.array.size != 5) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize_msgpack AllocatedBuffer invalid array " - "size: expected 5, got {}", - obj.via.array.size))); - } - - auto *array_items = obj.via.array.ptr; - - // Deserialize basic properties - // std::string segment_name = array_items[0].as(); - auto size = static_cast(array_items[0].as()); - void *buffer_ptr = reinterpret_cast(array_items[1].as()); - // auto status = static_cast(array_items[3].as()); - - // Get segment_id and find corresponding allocator - std::string segment_id = array_items[2].as(); - UUID segment_uuid; - bool success = StringToUuid(segment_id, segment_uuid); - if (!success) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize_msgpack AllocatedBuffer invalid segment " - "ID format: {}", - segment_id))); - } - - MountedSegment mountedSegment; - ErrorCode ret = - segment_view.GetMountedSegment(segment_uuid, mountedSegment); - if (ret != ErrorCode::OK) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format( - "deserialize_msgpack AllocatedBuffer " - "segment_view.GetMountedSegment() fail ret={},segment_id={}", - static_cast(ret), segment_id))); - } - - if (mountedSegment.status != SegmentStatus::OK) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize_msgpack AllocatedBuffer " - "mountedSegment.status!=OK status={} segment_id={}", - static_cast(mountedSegment.status), - segment_id))); - } - - std::shared_ptr allocator = - mountedSegment.buf_allocator; - // Check if allocator is valid - if (!allocator) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize_msgpack AllocatedBuffer invalid allocator " - "for segment {}", - segment_id))); - } - - // Deserialize offset_handle_ (if exists) - std::optional offsetHandle = - std::nullopt; - - bool has_offset_handle = array_items[3].as(); - if (has_offset_handle) { - auto offset_allocator = - std::dynamic_pointer_cast(allocator); - if (offset_allocator) { - // Use OffsetBufferAllocator's offset_allocator_ to create - // OffsetAllocationHandle - auto handle_result = - Serializer:: - deserialize(array_items[4], - offset_allocator->getOffsetAllocator()); - if (!handle_result) { - return tl::unexpected(handle_result.error()); - } - offsetHandle = std::move(*handle_result.value()); - } - } - - // Create AllocatedBuffer object - auto buffer = std::make_unique(allocator, buffer_ptr, size, - std::move(offsetHandle)); - if (mountedSegment.allocator_registration) { - mountedSegment.allocator_registration->BindBuffer(*buffer); - } - // buffer->status = status; - - return buffer; -} - -tl::expected Serializer::serialize( - const Replica &replica, const SegmentView &segment_view, - MsgpackPacker &packer) { - // Use unified array structure to pack Replica - // Format: [id(uint64), status(int16), replica_type(int8), payload] - packer.pack_array(4); - - // 1. Serialize id_ member variable - packer.pack(static_cast(replica.id_)); - - // 2. Serialize status_ member variable - packer.pack(static_cast(replica.status_)); - - // 3. Serialize replica type - auto replica_type = replica.type(); - packer.pack(static_cast(replica_type)); - - // 4. Serialize specific data by type - switch (replica_type) { - case ReplicaType::MEMORY: { - const auto *mem_data = - std::get_if(&replica.data_); - if (!mem_data || !mem_data->buffer) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("serialize_msgpack Replica memory buffer_ptr " - "is nullptr"))); - } - auto result = Serializer::serialize( - *mem_data->buffer, segment_view, packer); - if (!result) { - return tl::unexpected(result.error()); - } - break; - } - case ReplicaType::DISK: { - const auto *disk_data = - std::get_if(&replica.data_); - if (!disk_data) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "serialize_msgpack Replica missing DiskReplicaData")); - } - // Format: [file_path, object_size] - packer.pack_array(2); - packer.pack(disk_data->file_path); - packer.pack(static_cast(disk_data->object_size)); - break; - } - case ReplicaType::LOCAL_DISK: { - const auto *local_data = - std::get_if(&replica.data_); - if (!local_data) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "serialize_msgpack Replica missing LocalDiskReplicaData")); - } - // Format: [client_id_str, object_size, transport_endpoint] - packer.pack_array(3); - packer.pack(UuidToString(local_data->client_id)); - packer.pack(static_cast(local_data->object_size)); - packer.pack(local_data->transport_endpoint); - break; - } - case ReplicaType::DFS: { - const auto *dfs_data = std::get_if(&replica.data_); - if (!dfs_data) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "serialize_msgpack Replica missing DfsReplicaData")); - } - // Format: [file_path, offset, object_size, aligned_size, shard_idx] - packer.pack_array(5); - packer.pack(dfs_data->descriptor.file_path); - packer.pack(static_cast(dfs_data->descriptor.offset)); - packer.pack( - static_cast(dfs_data->descriptor.object_size)); - packer.pack( - static_cast(dfs_data->descriptor.aligned_size)); - packer.pack(static_cast(dfs_data->descriptor.shard_idx)); - break; - } - default: - // Unsupported replica type - packer.pack(static_cast(255)); - packer.pack_nil(); - break; - } - - return {}; -} - -auto Serializer::deserialize(const msgpack::object &obj, - const SegmentView &segment_view) - -> tl::expected { - // Check if object type is array (consistent with serialize) - if (obj.type != msgpack::type::ARRAY) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize_msgpack Replica invalid msgpack " - "data, expected array")); - } - - // Verify array size is correct (should have 4 elements: id, status, - // replica_type, payload) - if (obj.via.array.size != 4) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize_msgpack Replica invalid array size: " - "expected 4, got {}", - obj.via.array.size))); - } - - auto *array_items = obj.via.array.ptr; - - // 1. Deserialize id_ member variable - auto id = static_cast(array_items[0].as()); - - // 2. Deserialize status_ member variable - auto status = static_cast(array_items[1].as()); - - // 3. Deserialize replica_type - auto replica_type_code = array_items[2].as(); - - // 4. Parse payload by type - std::shared_ptr replica; - switch (replica_type_code) { - case static_cast(ReplicaType::MEMORY): { - // MEMORY: payload is AllocatedBuffer - auto buffer_result = Serializer::deserialize( - array_items[3], segment_view); - if (!buffer_result) { - return tl::unexpected(buffer_result.error()); - } - replica = std::make_shared( - std::move(buffer_result.value()), status); - break; - } - case static_cast(ReplicaType::DISK): { - const auto &payload = array_items[3]; - if (payload.type != msgpack::type::ARRAY || - payload.via.array.size != 2) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize_msgpack Replica DISK " - "payload is not valid array[2]")); - } - auto *payload_items = payload.via.array.ptr; - std::string file_path = payload_items[0].as(); - uint64_t object_size = payload_items[1].as(); - - replica = std::make_shared(std::move(file_path), - object_size, status); - break; - } - case static_cast(ReplicaType::LOCAL_DISK): { - const auto &payload = array_items[3]; - if (payload.type != msgpack::type::ARRAY || - payload.via.array.size != 3) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize_msgpack Replica LOCAL_DISK " - "payload is not valid array[3]")); - } - auto *payload_items = payload.via.array.ptr; - std::string client_id_str = payload_items[0].as(); - uint64_t object_size = payload_items[1].as(); - std::string transport_endpoint = payload_items[2].as(); - - UUID client_id; - if (!StringToUuid(client_id_str, client_id)) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize_msgpack Replica invalid client_id " - "UUID: {}", - client_id_str))); - } - - replica = std::make_shared( - client_id, object_size, std::move(transport_endpoint), status); - break; - } - case static_cast(ReplicaType::DFS): { - const auto &payload = array_items[3]; - if (payload.type != msgpack::type::ARRAY || - payload.via.array.size != 5) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize_msgpack Replica DFS " - "payload is not valid array[5]")); - } - auto *payload_items = payload.via.array.ptr; - DistributedFSDescriptor descriptor; - descriptor.file_path = payload_items[0].as(); - descriptor.offset = payload_items[1].as(); - descriptor.object_size = payload_items[2].as(); - descriptor.aligned_size = payload_items[3].as(); - descriptor.shard_idx = payload_items[4].as(); - - replica = std::make_shared(std::move(descriptor), status); - break; - } - default: - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize Replica invalid replica type: {}", - replica_type_code))); - } - - // Restore the original id (overwrite the auto-generated one) - // Note: refcnt_ is not restored, it remains 0 (default value) - replica->id_ = id; - - return replica; -} - -tl::expected Serializer::serialize( - const MountedSegment &mounted_segment, MsgpackPacker &packer) { - // Use array structure for packing, more efficient - // Format: [segment_id, segment_name, segment_base, segment_size, - // te_endpoint, status, has_buffer_allocator, buffer_allocator_data, - // host_id] - - packer.pack_array(9); - - // Serialize Segment info - packer.pack(UuidToString(mounted_segment.segment.id)); - packer.pack(mounted_segment.segment.name); - packer.pack(static_cast(mounted_segment.segment.base)); - packer.pack(static_cast(mounted_segment.segment.size)); - packer.pack(mounted_segment.segment.te_endpoint); - - // Serialize SegmentStatus - packer.pack(static_cast(mounted_segment.status)); - - // Serialize BufferAllocator - if (mounted_segment.buf_allocator) { - auto offsetAllocator = std::dynamic_pointer_cast( - mounted_segment.buf_allocator); - if (offsetAllocator) { - packer.pack(true); // Mark buffer allocator exists - auto result = Serializer::serialize( - *offsetAllocator, packer); - if (!result) { - return tl::unexpected(result.error()); - } - packer.pack(mounted_segment.segment.host_id); - return {}; - } - } - - packer.pack(false); // Mark no valid buffer allocator exists - packer.pack_nil(); - packer.pack(mounted_segment.segment.host_id); - return {}; -} - -tl::expected -Serializer::deserialize(const msgpack::object &obj) { - if (obj.type != msgpack::type::ARRAY) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize MountedSegment invalid serialized " - "state: not a msgpack array")); - } - - if (obj.via.array.size < 8) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize MountedSegment invalid array size")); - } - - MountedSegment mounted_segment; - msgpack::object *array = obj.via.array.ptr; - - try { - // Deserialize Segment info - std::string segment_id_str = array[0].as(); - UUID segment_uuid; - if (!StringToUuid(segment_id_str, segment_uuid)) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize MountedSegment invalid UUID {}", - segment_id_str))); - } - - mounted_segment.segment.id = segment_uuid; - mounted_segment.segment.name = array[1].as(); - mounted_segment.segment.base = - static_cast(array[2].as()); - mounted_segment.segment.size = - static_cast(array[3].as()); - mounted_segment.segment.te_endpoint = array[4].as(); - - // Deserialize SegmentStatus - mounted_segment.status = - static_cast(array[5].as()); - - // Deserialize BufferAllocator - bool has_buffer_allocator = array[6].as(); - if (has_buffer_allocator) { - auto allocatorResult = - Serializer::deserialize(array[7]); - if (allocatorResult) { - mounted_segment.buf_allocator = - std::move(allocatorResult.value()); - } else { - return tl::unexpected(allocatorResult.error()); - } - } - if (obj.via.array.size >= 9) { - mounted_segment.segment.host_id = array[8].as(); - } - } catch (const std::exception &e) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize MountedSegment failed: {}", e.what()))); - } - - return mounted_segment; -} - -tl::expected -Serializer::serialize( - const OffsetBufferAllocator &allocator, MsgpackPacker &packer) { - // Use array structure to pack OffsetBufferAllocator - // Format: [segment_name, base, total_size, current_size, - // transport_endpoint, offset_allocator] - - packer.pack_array(6); - - // Serialize basic properties - packer.pack(allocator.getSegmentName()); - packer.pack(static_cast(allocator.base())); - packer.pack(static_cast(allocator.capacity())); - packer.pack(static_cast(allocator.size())); - packer.pack(allocator.getTransportEndpoint()); - - // Serialize offset_allocator - return Serializer::serialize( - *allocator.getOffsetAllocator(), packer); -} - -auto Serializer::deserialize(const msgpack::object &obj) - -> tl::expected { - // Validate input state - if (obj.type != msgpack::type::ARRAY) { - return tl::unexpected( - SerializationError(ErrorCode::DESERIALIZE_FAIL, - "deserialize OffsetBufferAllocator invalid " - "serialized state: not a msgpack array")); - } - - if (obj.via.array.size != 6) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize OffsetBufferAllocator invalid array size")); - } - - try { - auto snapshot = ha::AllocatorSnapshotCodec::Decode(obj); - if (!snapshot) return tl::unexpected(snapshot.error()); - auto allocator = OffsetBufferAllocator::Restore(std::move(*snapshot)); - if (!allocator) { - return tl::unexpected(SerializationError( - allocator.error(), - "deserialize OffsetBufferAllocator rejected its state")); - } - return std::move(*allocator); - } catch (const std::exception &e) { - return tl::unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - fmt::format("deserialize OffsetBufferAllocator failed: {}", - e.what()))); - } -} - -} // namespace mooncake diff --git a/mooncake-store/src/task_manager.cpp b/mooncake-store/src/task_manager.cpp index 9ae2657045..4adba6926b 100644 --- a/mooncake-store/src/task_manager.cpp +++ b/mooncake-store/src/task_manager.cpp @@ -1,7 +1,5 @@ #include "task_manager.h" #include -#include -#include "common/zstd_util.h" namespace mooncake { @@ -258,33 +256,6 @@ void ScopedTaskWriteAccess::prune_expired_tasks() { } } -void ScopedTaskWriteAccess::restore_task(Task&& task) { - const UUID task_id = task.id; - manager_->all_tasks_[task_id] = std::move(task); - Task& stored_task = manager_->all_tasks_[task_id]; - - const UUID client_id = stored_task.assigned_client; - const TaskStatus status = stored_task.status; - - switch (status) { - case TaskStatus::PENDING: { - manager_->pending_tasks_[client_id].push(task_id); - manager_->total_pending_tasks_++; - break; - } - case TaskStatus::PROCESSING: { - manager_->processing_tasks_[client_id].insert(task_id); - manager_->total_processing_tasks_++; - break; - } - case TaskStatus::FAILED: - case TaskStatus::SUCCESS: { - manager_->finished_task_history_.push_back(task_id); - break; - } - } -} - ErrorCode ScopedTaskWriteAccess::fail_task_if_pending( const UUID& task_id, const std::string& message) { auto it = manager_->all_tasks_.find(task_id); @@ -305,164 +276,4 @@ ErrorCode ScopedTaskWriteAccess::fail_task_if_pending( return ErrorCode::OK; } -void ScopedTaskWriteAccess::clear_all() { - manager_->all_tasks_.clear(); - manager_->pending_tasks_.clear(); - manager_->processing_tasks_.clear(); - manager_->finished_task_history_.clear(); - manager_->total_pending_tasks_ = 0; - manager_->total_processing_tasks_ = 0; -} - -tl::expected, SerializationError> -TaskManagerSerializer::Serialize() { - if (!task_manager_) { - return tl::make_unexpected( - SerializationError(ErrorCode::SERIALIZE_FAIL, - "serialize TaskManager task_manager_ is null")); - } - - auto read_access = task_manager_->get_read_access(); - - // Collect and sort task_ids for deterministic serialization order - std::vector sorted_task_ids; - sorted_task_ids.reserve(read_access.size()); - for (const auto& [task_id, _] : read_access) { - sorted_task_ids.push_back(task_id); - } - std::sort(sorted_task_ids.begin(), sorted_task_ids.end()); - - msgpack::sbuffer sbuf; - msgpack::packer packer(&sbuf); - - packer.pack_array(read_access.size()); - - for (const auto& task_id : sorted_task_ids) { - auto task_opt = read_access.find_task_by_id(task_id); - if (!task_opt) continue; - const auto& task = *task_opt; - // Serialize Task as array: [id, type, status, payload, created_at, - // last_updated_at, message, assigned_client] - packer.pack_array(TaskManagerSerializer::kTaskSerializedFields); - packer.pack(UuidToString(task.id)); - packer.pack(static_cast(task.type)); - packer.pack(static_cast(task.status)); - packer.pack(task.payload); - packer.pack(static_cast( - std::chrono::duration_cast( - task.created_at.time_since_epoch()) - .count())); - packer.pack(static_cast( - std::chrono::duration_cast( - task.last_updated_at.time_since_epoch()) - .count())); - packer.pack(task.message); - packer.pack(UuidToString(task.assigned_client)); - } - - // Compress entire data - std::vector compressed_data = zstd_compress( - reinterpret_cast(sbuf.data()), sbuf.size(), 3); - - // Return compressed data - return compressed_data; -} - -tl::expected TaskManagerSerializer::Deserialize( - const std::vector& data) { - if (!task_manager_) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "deserialize TaskManager task_manager_ is null")); - } - - // Decompress data - std::vector decompressed_data; - try { - decompressed_data = zstd_decompress( - reinterpret_cast(data.data()), data.size(), - TaskManagerSerializer::kMaxDecompressedSize); - } catch (const std::exception& e) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "failed to decompress MessagePack data: " + std::string(e.what()))); - } - - // Parse MessagePack data - msgpack::object_handle oh; - try { - oh = msgpack::unpack( - reinterpret_cast(decompressed_data.data()), - decompressed_data.size()); - } catch (const std::exception& e) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "failed to unpack msgpack: " + std::string(e.what()))); - } - - const msgpack::object& obj = oh.get(); - if (obj.type != msgpack::type::ARRAY) { - return tl::make_unexpected(SerializationError( - ErrorCode::DESERIALIZE_FAIL, - "invalid MessagePack format: expected array of tasks")); - } - - auto write_access = task_manager_->get_write_access(); - - // Clear existing tasks - write_access.clear_all(); - - // Deserialize each task - for (uint32_t i = 0; i < obj.via.array.size; ++i) { - const msgpack::object& task_obj = obj.via.array.ptr[i]; - if (task_obj.type != msgpack::type::ARRAY || - task_obj.via.array.size != - TaskManagerSerializer::kTaskSerializedFields) { - LOG(WARNING) << "Invalid task format for task index " << i - << ", skipping deserialization"; - continue; - } - - const msgpack::object* arr = task_obj.via.array.ptr; - - Task task; - std::string id_str = arr[0].as(); - if (!StringToUuid(id_str, task.id)) { - LOG(WARNING) << "Invalid UUID format for task id " << id_str - << ", skipping deserialization"; - continue; - } - try { - task.type = static_cast(arr[1].as()); - task.status = static_cast(arr[2].as()); - task.payload = arr[3].as(); - task.created_at = std::chrono::system_clock::time_point( - std::chrono::milliseconds(arr[4].as())); - task.last_updated_at = std::chrono::system_clock::time_point( - std::chrono::milliseconds(arr[5].as())); - task.message = arr[6].as(); - } catch (const std::exception& e) { - LOG(WARNING) << "Invalid task field types for task index " << i - << ": " << e.what() << ", skipping deserialization"; - continue; - } - - std::string assigned_str = arr[7].as(); - if (!StringToUuid(assigned_str, task.assigned_client)) { - LOG(WARNING) << "Invalid UUID format for assigned_client " - << assigned_str << ", skipping deserialization"; - continue; - } - write_access.restore_task(std::move(task)); - } - return {}; -} - -void TaskManagerSerializer::Reset() { - if (!task_manager_) { - return; - } - auto write_access = task_manager_->get_write_access(); - write_access.clear_all(); -} } // namespace mooncake diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 5dab2557c4..726af40629 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -236,9 +236,6 @@ target_include_directories(local_ssd_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) target_link_libraries(local_ssd_test PRIVATE mooncake_local_ssd gtest gtest_main) add_test(NAME local_ssd_test COMMAND local_ssd_test) -add_store_test(local_ssd_codec_test ha/snapshot/local_ssd_codec_test.cpp) -add_store_test(store_resource_snapshot_codec_test - ha/snapshot/store_resource_snapshot_codec_test.cpp) add_store_test(offset_allocator_test offset_allocator_test.cpp) add_store_test(byte_size_test byte_size_test.cpp) add_store_test(network_test network_test.cpp) @@ -298,7 +295,6 @@ set_tests_properties(client_metric_cluster_id_test PROPERTIES ENVIRONMENT "MC_STORE_CLUSTER_ID=test-cluster") add_store_test(ssd_metrics_test ssd_metrics_test.cpp) add_store_test(allocator_metric_test allocator_metric_test.cpp) -add_store_test(serializer_test serializer_test.cpp) add_store_test(dfs_posix_test dfs_posix_test.cpp) add_store_test(dfs_immutable_bucket_allocator_test dfs_immutable_bucket_allocator_test.cpp) diff --git a/mooncake-store/tests/buffer_allocator_test.cpp b/mooncake-store/tests/buffer_allocator_test.cpp index 605d51c995..124204a2c1 100644 --- a/mooncake-store/tests/buffer_allocator_test.cpp +++ b/mooncake-store/tests/buffer_allocator_test.cpp @@ -284,124 +284,6 @@ TEST_F(BufferAllocatorTest, CachelibCreateRejectsInvalidMemoryLayout) { } } -TEST_F(BufferAllocatorTest, OffsetRestoreFactoryRejectsInvalidState) { - auto& metrics = MasterMetricManager::instance(); - const auto baseline = metrics.get_allocated_mem_size(); - const std::string segment_name = "restore-factory"; - constexpr size_t kBase = 0x100000000ULL; - constexpr size_t kCapacity = 16U * 1024 * 1024; - auto state = offset_allocator::OffsetAllocator::create(kBase, kCapacity); - ASSERT_NE(state, nullptr); - - auto snapshot = [&] { - return OffsetBufferAllocatorSnapshot{segment_name, - kBase, - kCapacity, - 0, - segment_name + "-endpoint", - state->CaptureSnapshot()}; - }; - const auto expect_rejected = [](OffsetBufferAllocatorSnapshot candidate) { - auto restored = OffsetBufferAllocator::Restore(std::move(candidate)); - ASSERT_FALSE(restored.has_value()); - EXPECT_EQ(restored.error(), ErrorCode::INVALID_PARAMS); - }; - auto missing_layout = snapshot(); - missing_layout.allocation_state.layout.reset(); - expect_rejected(std::move(missing_layout)); - auto invalid_usage = snapshot(); - invalid_usage.used_bytes = kCapacity + 1; - expect_rejected(std::move(invalid_usage)); - auto mismatched_base = snapshot(); - ++mismatched_base.allocation_state.base; - expect_rejected(std::move(mismatched_base)); - auto overflow = snapshot(); - overflow.used_bytes = std::numeric_limits::max(); - expect_rejected(std::move(overflow)); - // The persisted counters must describe the persisted layout. - auto wrong_alloc_num = snapshot(); - ++wrong_alloc_num.allocation_state.allocated_num; - expect_rejected(std::move(wrong_alloc_num)); - auto wrong_alloc_size = snapshot(); - wrong_alloc_size.allocation_state.allocated_size = kCapacity; - expect_rejected(std::move(wrong_alloc_size)); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline); - - { - auto restored = OffsetBufferAllocator::Restore(snapshot()); - ASSERT_TRUE(restored.has_value()); - EXPECT_EQ((*restored)->size(), 0U); - EXPECT_EQ((*restored)->capacity(), kCapacity); - - auto buffer = (*restored)->allocate(4096); - ASSERT_NE(buffer, nullptr); - EXPECT_EQ((*restored)->size(), 4096U); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline + 4096); - } - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline); -} - -TEST_F(BufferAllocatorTest, - SnapshotValidationReportsErrorsWithoutPublishingUsage) { - auto& metrics = MasterMetricManager::instance(); - const auto baseline = metrics.get_allocated_mem_size(); - constexpr size_t kBase = 0x100000000ULL; - constexpr size_t kCapacity = 16U << 20; - auto state = - offset_allocator::OffsetAllocator::create(kBase, kCapacity, 16, 32); - auto allocation = state->allocate(4097); - ASSERT_TRUE(allocation.has_value()); - auto snapshot = [&] { - return OffsetBufferAllocatorSnapshot{ - .segment_name = "validate-snapshot", - .base = kBase, - .capacity = kCapacity, - .used_bytes = 4097, - .transport_endpoint = "endpoint", - .allocation_state = state->CaptureSnapshot(), - }; - }; - struct Fault { - const char* diagnostic; - std::function apply; - }; - const std::vector faults{ - {"bounds", [](auto& s) { ++s.base; }}, - {"bounds", [](auto& s) { ++s.capacity; }}, - {"usage", [](auto& s) { s.used_bytes = s.capacity + 1; }}, - {"used_bytes", [](auto& s) { s.used_bytes = 0; }}, - {"layout", [](auto& s) { s.allocation_state.layout.reset(); }}, - {"multiplier_bits", - [](auto& s) { s.allocation_state.multiplier_bits = 64; }}, - {"allocated_num", [](auto& s) { ++s.allocation_state.allocated_num; }}, - {"allocated_size", - [](auto& s) { s.used_bytes = s.allocation_state.allocated_size = 0; }}, - }; - for (const auto& fault : faults) { - SCOPED_TRACE(fault.diagnostic); - auto candidate = snapshot(); - fault.apply(candidate); - const auto valid = candidate.Validate(); - ASSERT_FALSE(valid.has_value()); - EXPECT_NE(valid.error().find(fault.diagnostic), std::string::npos); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline); - auto restored = OffsetBufferAllocator::Restore(std::move(candidate)); - ASSERT_FALSE(restored.has_value()); - EXPECT_EQ(restored.error(), ErrorCode::INVALID_PARAMS); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline); - } - auto clean = snapshot(); - ASSERT_TRUE(clean.Validate().has_value()); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline); - { - auto restored = OffsetBufferAllocator::Restore(std::move(clean)); - ASSERT_TRUE(restored.has_value()); - EXPECT_EQ((*restored)->size(), 4097U); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline + 4097); - } - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline); -} - TEST_F(BufferAllocatorTest, ImportCachelibAllocationsAtOriginalAddresses) { constexpr uintptr_t kBase = 0x1C0000000ULL; constexpr size_t kCapacity = 4 * facebook::cachelib::Slab::kSize; diff --git a/mooncake-store/tests/ha/snapshot/local_ssd_codec_test.cpp b/mooncake-store/tests/ha/snapshot/local_ssd_codec_test.cpp deleted file mode 100644 index 2bbf6e2b08..0000000000 --- a/mooncake-store/tests/ha/snapshot/local_ssd_codec_test.cpp +++ /dev/null @@ -1,79 +0,0 @@ -#include - -#include "ha/snapshot/local_ssd_codec.h" - -namespace mooncake::ha { -namespace { - -OffloadTaskItem Offload(std::string tenant, std::string key, int64_t size) { - return OffloadTaskItem{ - .tenant_id = std::move(tenant), .key = std::move(key), .size = size}; -} - -std::string PackState(const LocalSsdPersistedState& state) { - msgpack::sbuffer buffer; - MsgpackPacker packer(&buffer); - EXPECT_TRUE(LocalSsdCodec::Encode(state, packer).has_value()); - return std::string(buffer.data(), buffer.size()); -} - -TEST(LocalSsdCodecTest, RoundTripIsDeterministic) { - LocalSsdPersistedState state; - state[{9, 2}] = LocalSsdPersistedClient{ - .enable_offloading = false, - .total_capacity_bytes = 10, - .pending_offloads = {{"z", Offload("tenant", "z", 1)}}}; - state[{1, 8}] = LocalSsdPersistedClient{ - .enable_offloading = true, - .total_capacity_bytes = 20, - .pending_offloads = {{"b", Offload("tenant", "b", 2)}, - {"a", Offload("tenant", "a", 3)}}}; - - auto first = PackState(state); - auto second = PackState(state); - EXPECT_EQ(first, second); - auto object = msgpack::unpack(first.data(), first.size()); - auto decoded = LocalSsdCodec::Decode(&object.get()); - ASSERT_TRUE(decoded.has_value()); - EXPECT_EQ(PackState(*decoded), first); -} - -TEST(LocalSsdCodecTest, DecodesLegacySizeAndPreservesEncodedKey) { - const std::string encoded_key("legacy\0key", 10); - msgpack::sbuffer buffer; - MsgpackPacker packer(&buffer); - packer.pack_map(1); - packer.pack(UuidToString(UUID{11, 12})); - packer.pack_array(4); - packer.pack(true); - packer.pack(uint64_t{1}); - packer.pack(encoded_key); - packer.pack(int64_t{99}); - - auto object = msgpack::unpack(buffer.data(), buffer.size()); - auto decoded = LocalSsdCodec::Decode(&object.get()); - ASSERT_TRUE(decoded.has_value()); - const auto& client = decoded->at(UUID{11, 12}); - EXPECT_EQ(client.total_capacity_bytes, 0); - ASSERT_EQ(client.pending_offloads.size(), 1); - auto task = client.pending_offloads.find(encoded_key); - ASSERT_NE(task, client.pending_offloads.end()); - EXPECT_EQ(task->second.tenant_id, "legacy"); - EXPECT_EQ(task->second.key, "key"); - EXPECT_EQ(task->second.size, 99); -} - -TEST(LocalSsdCodecTest, MissingFieldIsEmptyAndInvalidInputFails) { - auto missing = LocalSsdCodec::Decode(nullptr); - ASSERT_TRUE(missing.has_value()); - EXPECT_TRUE(missing->empty()); - - msgpack::sbuffer buffer; - MsgpackPacker packer(&buffer); - packer.pack_array(0); - auto object = msgpack::unpack(buffer.data(), buffer.size()); - EXPECT_FALSE(LocalSsdCodec::Decode(&object.get()).has_value()); -} - -} // namespace -} // namespace mooncake::ha diff --git a/mooncake-store/tests/ha/snapshot/store_resource_snapshot_codec_test.cpp b/mooncake-store/tests/ha/snapshot/store_resource_snapshot_codec_test.cpp deleted file mode 100644 index 85d2eddf65..0000000000 --- a/mooncake-store/tests/ha/snapshot/store_resource_snapshot_codec_test.cpp +++ /dev/null @@ -1,881 +0,0 @@ -#include "ha/snapshot/store_resource_snapshot_codec.h" - -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -#include "common/zstd_util.h" -#include "ha/snapshot/local_ssd_codec.h" -#include "master_metric_manager.h" -#include "segment.h" -#include "segment/pool.h" -#include "segment/pool_read_access.h" -#include "segment/pool_write_access.h" -#include "segment/snapshot.h" -#include "serialize/serializer.h" - -namespace mooncake::ha { -namespace { - -constexpr size_t kRegionSize = 16U * 1024 * 1024; - -RegionDriverRegistry Drivers(bool enable_cxl = false) { - RegionDriverConfig config; - config.memory_allocator = BufferAllocatorType::OFFSET; - if (enable_cxl) { - config.cxl = CxlRegionDriverConfig{"cxl-test-path", kRegionSize}; - } - auto drivers = CreateRegionDrivers(config); - EXPECT_TRUE(drivers.has_value()); - return std::move(drivers.value()); -} - -Segment MakeSegment(size_t index, std::string name, - std::string protocol = "tcp", std::string host = {}) { - Segment segment; - segment.id = generate_uuid(); - segment.name = std::move(name); - segment.base = 0x100000000ULL + index * 0x2000000ULL; - segment.size = kRegionSize; - segment.te_endpoint = segment.name + "-endpoint"; - segment.protocol = std::move(protocol); - segment.host_id = std::move(host); - return segment; -} - -void CommitUnmount(SegmentPool& pool, const Segment& segment, - const UUID& client_id) { - auto transaction = [&] { - auto access = pool.AcquireWriteAccess(); - return access.PrepareUnmount(segment.id, client_id); - }(); - ASSERT_TRUE(transaction.has_value()); - auto access = pool.AcquireWriteAccess(); - ASSERT_EQ(std::move(*transaction).Commit(access), ErrorCode::OK); - EXPECT_EQ(std::move(*transaction).Commit(access), - ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); - EXPECT_EQ(std::move(*transaction).Rollback(access), - ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); -} - -tl::expected, SerializationError> CaptureAndEncode( - const SegmentPool& pool, const LocalSsdPersistedState& local_ssd) { - auto snapshot = pool.CaptureSnapshot(); - if (!snapshot) { - return tl::make_unexpected( - SerializationError(snapshot.error(), "capture SegmentPool failed")); - } - return StoreResourceSnapshotCodec::Encode(*snapshot, local_ssd); -} - -tl::expected DecodeAndRestore( - SegmentPool& pool, const std::vector& bytes, - bool account_capacity) { - auto decoded = StoreResourceSnapshotCodec::Decode(bytes); - if (!decoded) return tl::make_unexpected(decoded.error()); - auto restored = pool.RestoreSnapshot(std::move(decoded->segment_pool), - account_capacity); - if (!restored) - return tl::make_unexpected( - SerializationError(restored.error(), "restore SegmentPool failed")); - return std::move(decoded->local_ssd); -} - -// Mutate the unpacked wire payload rather than duplicating the codec in tests. -std::vector RewriteSnapshot( - const std::vector& snapshot, - const std::function& rewrite) { - const auto bytes = zstd_decompress(snapshot); - auto unpacked = msgpack::unpack(reinterpret_cast(bytes.data()), - bytes.size()); - auto root = unpacked.get(); - rewrite(root, *unpacked.zone()); - - msgpack::sbuffer buffer; - msgpack::pack(buffer, root); - return zstd_compress(reinterpret_cast(buffer.data()), - buffer.size(), 3); -} - -msgpack::object& SnapshotField(msgpack::object& root, std::string_view name) { - for (uint32_t i = 0; i < root.via.map.size; ++i) { - auto& item = root.via.map.ptr[i]; - if (item.key.as() == name) { - return item.val; - } - } - - throw std::runtime_error("missing snapshot field"); -} - -std::string PackLocalSsd(const LocalSsdPersistedState& state) { - msgpack::sbuffer buffer; - MsgpackPacker packer(&buffer); - EXPECT_TRUE(LocalSsdCodec::Encode(state, packer).has_value()); - return {buffer.data(), buffer.size()}; -} - -std::vector ReplaceSnapshotActiveNames( - const std::vector& snapshot, - const std::vector& active_names) { - return RewriteSnapshot(snapshot, [&](auto& root, auto& zone) { - SnapshotField(root, "an") = msgpack::object(active_names, zone); - }); -} - -std::vector ReplaceSnapshotAllocatorType( - const std::vector& snapshot, uint64_t allocator_type) { - return RewriteSnapshot(snapshot, [&](auto& root, auto&) { - SnapshotField(root, "ma") = msgpack::object(allocator_type); - }); -} - -} // namespace - -TEST(StoreResourceSnapshotCodecTest, - CapturedSnapshotOutlivesPoolAndPreservesOriginalLayout) { - auto& metrics = MasterMetricManager::instance(); - const auto baseline_used = metrics.get_allocated_mem_size(); - const auto baseline_capacity = metrics.get_total_mem_capacity(); - const auto segment = MakeSegment(0, "detached-snapshot", "tcp", "host-a"); - const UUID client = generate_uuid(); - std::optional snapshot; - std::vector original_bytes; - uintptr_t first_address = 0; - uintptr_t last_address = 0; - { - SegmentPool pool(Drivers()); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, client), - ErrorCode::OK); - auto allocator = pool.AcquireReadAccess().GetAllocator(segment.id); - auto first = allocator->allocate(4096); - auto hole = allocator->allocate(4096); - auto last = allocator->allocate(4096); - ASSERT_NE(first, nullptr); - ASSERT_NE(hole, nullptr); - ASSERT_NE(last, nullptr); - first_address = reinterpret_cast(first->data()); - last_address = reinterpret_cast(last->data()); - hole.reset(); - const auto used = metrics.get_allocated_mem_size(); - auto captured = pool.CaptureSnapshot(); - ASSERT_TRUE(captured.has_value()); - snapshot.emplace(std::move(*captured)); - EXPECT_EQ(metrics.get_allocated_mem_size(), used); - auto encoded = StoreResourceSnapshotCodec::Encode(*snapshot, {}); - ASSERT_TRUE(encoded.has_value()); - original_bytes = std::move(*encoded); - - // Change both allocation bins/nodes and catalog/placement after - // capture. - first.reset(); - last.reset(); - auto all = allocator->allocate(kRegionSize); - ASSERT_NE(all, nullptr); - ASSERT_EQ(pool.AcquireWriteAccess().SetSegmentStatusByName( - segment.name, SegmentStatus::DRAINED), - ErrorCode::OK); - pool.AcquireWriteAccess().Clear(); - } - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline_used); - EXPECT_EQ(metrics.get_total_mem_capacity(), baseline_capacity); - auto encoded = StoreResourceSnapshotCodec::Encode(*snapshot, {}); - ASSERT_TRUE(encoded.has_value()); - EXPECT_EQ(*encoded, original_bytes); - snapshot.reset(); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline_used); - - SegmentPool restored(Drivers()); - ASSERT_TRUE(DecodeAndRestore(restored, *encoded, false).has_value()); - EXPECT_EQ(restored.GetMemoryUsage().used_bytes, 8192U); - auto access = restored.AcquireReadAccess(); - const auto* mounted = access.Catalog().Find(segment.id); - ASSERT_NE(mounted, nullptr); - EXPECT_EQ(mounted->status, SegmentStatus::OK); - EXPECT_EQ(mounted->client_id, client); - EXPECT_EQ(mounted->segment.host_id, segment.host_id); - EXPECT_TRUE(access.Placement().Contains(segment.name, - AllocationCandidateKind::NATIVE)); - auto next = access.GetAllocator(segment.id)->allocate(4096); - ASSERT_NE(next, nullptr); - EXPECT_NE(reinterpret_cast(next->data()), first_address); - EXPECT_NE(reinterpret_cast(next->data()), last_address); -} - -TEST(StoreResourceSnapshotCodecTest, - DecodeReturnsDetachedStateWithoutMetricSideEffects) { - SegmentPool source(Drivers()); - const auto segment = MakeSegment(0, "decode-detached"); - const auto client = generate_uuid(); - ASSERT_EQ(source.AcquireWriteAccess().MountSegment(segment, client), - ErrorCode::OK); - auto allocator = source.AcquireReadAccess().GetAllocator(segment.id); - auto buffer = allocator->allocate(4096); - ASSERT_NE(buffer, nullptr); - auto encoded = CaptureAndEncode(source, {}); - ASSERT_TRUE(encoded); - auto& metrics = MasterMetricManager::instance(); - const auto used = metrics.get_allocated_mem_size(); - const auto capacity = metrics.get_total_mem_capacity(); - { - auto decoded = StoreResourceSnapshotCodec::Decode(*encoded); - ASSERT_TRUE(decoded); - ASSERT_EQ(decoded->segment_pool.regions.size(), 1U); - EXPECT_EQ(decoded->segment_pool.regions.front().allocator.used_bytes, - 4096U); - EXPECT_EQ(metrics.get_allocated_mem_size(), used); - EXPECT_EQ(metrics.get_total_mem_capacity(), capacity); - } - EXPECT_EQ(metrics.get_allocated_mem_size(), used); - EXPECT_EQ(metrics.get_total_mem_capacity(), capacity); -} - -TEST(StoreResourceSnapshotCodecTest, SnapshotDecodeReencodePreservesBytes) { - SegmentPool pool(Drivers()); - const UUID client{1, 1}; - auto first = MakeSegment(0, "roundtrip-z"); - auto second = MakeSegment(1, "roundtrip-a"); - auto inactive = MakeSegment(2, "roundtrip-inactive"); - // ID order differs from both mount order and active-name order. - first.id = UUID{0, 3}; - second.id = UUID{0, 1}; - inactive.id = UUID{0, 2}; - { - auto access = pool.AcquireWriteAccess(); - for (const auto* segment : {&first, &second, &inactive}) { - ASSERT_EQ(access.MountSegment(*segment, client), ErrorCode::OK); - } - ASSERT_EQ(access.SetSegmentStatusByName(inactive.name, - SegmentStatus::DRAINED), - ErrorCode::OK); - } - auto encoded = CaptureAndEncode(pool, {}); - ASSERT_TRUE(encoded.has_value()) << encoded.error().message; - - for (bool reverse_region_order : {false, true}) { - SCOPED_TRACE(reverse_region_order); - auto input = *encoded; - if (reverse_region_order) { - input = RewriteSnapshot(*encoded, [](auto& root, auto&) { - auto& regions = SnapshotField(root, "ms").via.map; - std::reverse(regions.ptr, regions.ptr + regions.size); - }); - } - auto decoded = StoreResourceSnapshotCodec::Decode(input); - ASSERT_TRUE(decoded.has_value()) << decoded.error().message; - const auto& snapshot = decoded->segment_pool; - ASSERT_EQ(snapshot.regions.size(), 3U); - EXPECT_EQ(snapshot.regions[0].mounted.segment.id, second.id); - EXPECT_EQ(snapshot.regions[1].mounted.segment.id, inactive.id); - EXPECT_EQ(snapshot.regions[2].mounted.segment.id, first.id); - EXPECT_EQ(snapshot.active_names, - (std::vector{first.name, second.name})); - - auto reencoded = - StoreResourceSnapshotCodec::Encode(snapshot, decoded->local_ssd); - ASSERT_TRUE(reencoded.has_value()) << reencoded.error().message; - EXPECT_EQ(*reencoded, *encoded); - } -} - -TEST(StoreResourceSnapshotCodecTest, - OutOfRangeSnapshotAllocatorTypeReturnsError) { - SegmentPool pool(Drivers()); - const UUID client = generate_uuid(); - auto segment = MakeSegment(0, "snapshot-invalid-allocator-type"); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, client), - ErrorCode::OK); - auto encoded = CaptureAndEncode(pool, LocalSsdPersistedState{}); - ASSERT_TRUE(encoded.has_value()) << encoded.error().message; - - for (const uint64_t value : - {uint64_t{std::numeric_limits::max()} + 1, - std::numeric_limits::max()}) { - auto corrupted = ReplaceSnapshotAllocatorType(*encoded, value); - ASSERT_NO_THROW({ - auto result = DecodeAndRestore(pool, corrupted, false); - ASSERT_FALSE(result.has_value()); - EXPECT_EQ(result.error().code, ErrorCode::DESERIALIZE_FAIL); - }); - EXPECT_NE(pool.AcquireReadAccess().Catalog().Find(segment.id), nullptr); - } - CommitUnmount(pool, segment, client); -} - -TEST(StoreResourceSnapshotCodecTest, - InvalidSnapshotDoesNotReplacePublishedPool) { - SegmentPool source(Drivers()); - const UUID client = generate_uuid(); - auto segment = MakeSegment(0, "snapshot-active"); - ASSERT_EQ(source.AcquireWriteAccess().MountSegment(segment, client), - ErrorCode::OK); - auto encoded = CaptureAndEncode(source, LocalSsdPersistedState{}); - ASSERT_TRUE(encoded.has_value()) << encoded.error().message; - - SegmentPool restored(Drivers()); - auto stale = MakeSegment(1, "published"); - ASSERT_EQ(restored.AcquireWriteAccess().MountSegment(stale, client), - ErrorCode::OK); - auto corrupted = ReplaceSnapshotActiveNames(*encoded, {}); - auto result = DecodeAndRestore(restored, corrupted, false); - ASSERT_FALSE(result.has_value()); - EXPECT_EQ(result.error().code, ErrorCode::DESERIALIZE_FAIL); - - { - auto view = restored.AcquireReadAccess(); - EXPECT_NE(view.Catalog().Find(stale.id), nullptr); - EXPECT_EQ(view.Catalog().Find(segment.id), nullptr); - } - - CommitUnmount(source, segment, client); -} - -TEST(StoreResourceSnapshotCodecTest, - SnapshotInteroperatesWithLegacySerializer) { - SegmentPool source(Drivers()); - const UUID client = generate_uuid(); - // Preserve insertion order, not lexicographic name order. - auto first = MakeSegment(0, "compat-z", "tcp", "host-a"); - auto second = MakeSegment(1, "compat-a", "tcp", "host-b"); - ASSERT_EQ(source.AcquireWriteAccess().MountSegment(first, client), - ErrorCode::OK); - ASSERT_EQ(source.AcquireWriteAccess().MountSegment(second, client), - ErrorCode::OK); - - LocalSsdPersistedState ssd; - ssd[client] = LocalSsdPersistedClient{ - .enable_offloading = true, - .total_capacity_bytes = 12345, - .pending_offloads = { - {"tenant/key", - OffloadTaskItem{ - .tenant_id = "tenant", .key = "key", .size = 4096}}}}; - auto encoded = CaptureAndEncode(source, ssd); - ASSERT_TRUE(encoded.has_value()); - - SegmentManager legacy(BufferAllocatorType::OFFSET); - SegmentSerializer serializer(&legacy); - auto old_decoded = serializer.Deserialize(*encoded); - ASSERT_TRUE(old_decoded.has_value()) << old_decoded.error().message; - EXPECT_EQ(PackLocalSsd(*old_decoded), PackLocalSsd(ssd)); - - MountedSegment mounted; - ASSERT_EQ(legacy.getView().GetMountedSegment(first.id, mounted), - ErrorCode::OK); - EXPECT_EQ(mounted.segment.host_id, first.host_id); - - auto old_encoded = serializer.Serialize(*old_decoded); - ASSERT_TRUE(old_encoded.has_value()); - - SegmentPool restored(Drivers()); - auto decoded = DecodeAndRestore(restored, *old_encoded, false); - ASSERT_TRUE(decoded.has_value()) << decoded.error().message; - EXPECT_EQ(PackLocalSsd(*decoded), PackLocalSsd(ssd)); - { - auto access = restored.AcquireReadAccess(); - ASSERT_EQ(access.Catalog().Regions().size(), 2U); - ASSERT_NE(access.Catalog().Find(first.id), nullptr); - EXPECT_EQ(access.Catalog().Find(first.id)->client_id, client); - std::vector names; - access.Placement().GetActiveSegmentNames( - AllocationCandidateKind::NATIVE, names); - EXPECT_EQ(names, (std::vector{first.name, second.name})); - } - - // Historical MountedSegment arrays had no host_id, and ld was optional. - auto historical = RewriteSnapshot(*old_encoded, [](auto& root, auto&) { - auto& regions = SnapshotField(root, "ms"); - for (uint32_t i = 0; i < regions.via.map.size; ++i) { - regions.via.map.ptr[i].val.via.array.size = 8; - } - - for (uint32_t i = 0; i < root.via.map.size; ++i) { - if (root.via.map.ptr[i].key.template as() == "ld") { - root.via.map.ptr[i] = root.via.map.ptr[root.via.map.size - 1]; - --root.via.map.size; - break; - } - } - }); - decoded = DecodeAndRestore(restored, historical, false); - ASSERT_TRUE(decoded.has_value()); - EXPECT_TRUE(decoded->empty()); - EXPECT_TRUE(restored.AcquireReadAccess() - .Catalog() - .Find(first.id) - ->segment.host_id.empty()); -} - -TEST(StoreResourceSnapshotCodecTest, - MalformedSnapshotsPreservePublishedResourcesAndMetrics) { - SegmentPool pool(Drivers()); - const UUID client = generate_uuid(); - const auto segment = MakeSegment(0, "snapshot-atomic"); - const auto other = MakeSegment(1, "snapshot-atomic-other"); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, client), - ErrorCode::OK); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(other, client), - ErrorCode::OK); - auto allocator = pool.AcquireReadAccess().GetAllocator(segment.id); - auto live = allocator->allocate(4096); - ASSERT_NE(live, nullptr); - auto encoded = CaptureAndEncode(pool, {}); - ASSERT_TRUE(encoded.has_value()); - auto& metrics = MasterMetricManager::instance(); - const auto used = metrics.get_allocated_mem_size(); - const auto capacity = metrics.get_total_mem_capacity(); - - using Rewrite = std::function; - const std::vector> cases{ - {"missing owner", - [](auto& root, auto&) { SnapshotField(root, "cs").via.map.size = 0; }}, - {"duplicate ownership", - [](auto& root, auto&) { - auto& ids = SnapshotField(root, "cs").via.map.ptr[0].val; - ids.via.array.ptr[1] = ids.via.array.ptr[0]; - }}, - {"unknown owned region", - [](auto& root, auto& zone) { - SnapshotField(root, "cs").via.map.ptr[0].val.via.array.ptr[0] = - msgpack::object(UuidToString(UUID{99, 99}), zone); - }}, - {"duplicate region", - [](auto& root, auto&) { - auto& regions = SnapshotField(root, "ms"); - regions.via.map.ptr[1] = regions.via.map.ptr[0]; - }}, - {"invalid status", - [](auto& root, auto&) { - SnapshotField(root, "ms").via.map.ptr[1].val.via.array.ptr[5] = - msgpack::object(99); - }}, - {"mismatched base", - [](auto& root, auto&) { - SnapshotField(root, "ms").via.map.ptr[1].val.via.array.ptr[2] = - msgpack::object(uint64_t{1}); - }}, - {"mismatched capacity", - [](auto& root, auto&) { - SnapshotField(root, "ms").via.map.ptr[1].val.via.array.ptr[3] = - msgpack::object(uint64_t{1}); - }}, - {"mismatched name", - [](auto& root, auto& zone) { - SnapshotField(root, "ms").via.map.ptr[1].val.via.array.ptr[1] = - msgpack::object(std::string("wrong"), zone); - }}, - {"mismatched endpoint", - [](auto& root, auto& zone) { - SnapshotField(root, "ms").via.map.ptr[1].val.via.array.ptr[4] = - msgpack::object(std::string("wrong"), zone); - }}, - {"broken allocator", - [](auto& root, auto&) { - SnapshotField(root, "ms").via.map.ptr[1].val.via.array.ptr[7] = - msgpack::object(); - }}, - {"usage exceeds capacity", - [](auto& root, auto&) { - // Mirror the new usage onto the inner allocated_size so the - // outer/inner consistency check passes and the bound itself is - // exercised. - auto& allocator = - SnapshotField(root, "ms").via.map.ptr[1].val.via.array.ptr[7]; - auto usage = msgpack::object(uint64_t{kRegionSize + 1}); - allocator.via.array.ptr[3] = usage; - allocator.via.array.ptr[5].via.array.ptr[3] = usage; - }}, - {"usage overflows metrics", - [](auto& root, auto&) { - auto& allocator = - SnapshotField(root, "ms").via.map.ptr[1].val.via.array.ptr[7]; - auto usage = msgpack::object(std::numeric_limits::max()); - allocator.via.array.ptr[3] = usage; - allocator.via.array.ptr[5].via.array.ptr[3] = usage; - }}, - {"second adoption fails", - [&](auto& root, auto& zone) { - // A zero UUID passes payload consistency checks but is rejected by - // the memory driver. The first active region has already staged. - const auto zero = msgpack::object(UuidToString(UUID{0, 0}), zone); - auto& regions = SnapshotField(root, "ms"); - for (uint32_t i = 0; i < regions.via.map.size; ++i) { - auto& item = regions.via.map.ptr[i]; - if (item.key.template as() == - UuidToString(other.id)) { - item.key = zero; - item.val.via.array.ptr[0] = zero; - } - } - auto& ids = SnapshotField(root, "cs").via.map.ptr[0].val; - for (uint32_t i = 0; i < ids.via.array.size; ++i) { - if (ids.via.array.ptr[i].template as() == - UuidToString(other.id)) { - ids.via.array.ptr[i] = zero; - } - } - }}, - {"invalid local ssd", - [](auto& root, auto&) { - SnapshotField(root, "ld") = msgpack::object(42); - }}, - }; - - for (const auto& [name, rewrite] : cases) { - SCOPED_TRACE(name); - auto corrupted = RewriteSnapshot(*encoded, rewrite); - auto result = DecodeAndRestore(pool, corrupted, true); - ASSERT_FALSE(result.has_value()); - EXPECT_EQ(result.error().code, name == "second adoption fails" - ? ErrorCode::INVALID_PARAMS - : ErrorCode::DESERIALIZE_FAIL); - EXPECT_EQ(pool.AcquireReadAccess().GetAllocator(segment.id), allocator); - - if (name == "second adoption fails") { - EXPECT_EQ(result.error().message, "restore SegmentPool failed"); - } else if (name == "usage exceeds capacity" || - name == "usage overflows metrics") { - // Decode propagates the detailed reason from the snapshot. - EXPECT_NE(result.error().message.find("invalid usage"), - std::string::npos) - << result.error().message; - } - - EXPECT_EQ(pool.AcquireReadAccess().Catalog().Regions().size(), 2U); - EXPECT_TRUE(pool.AcquireReadAccess().Placement().Contains( - segment.name, AllocationCandidateKind::NATIVE)); - EXPECT_EQ(pool.GetMemoryUsage().used_bytes, 4096U); - EXPECT_EQ(pool.GetMemoryUsage().capacity_bytes, 2 * kRegionSize); - EXPECT_EQ(metrics.get_allocated_mem_size(), used); - EXPECT_EQ(metrics.get_total_mem_capacity(), capacity); - EXPECT_EQ(metrics.get_segment_allocated_mem_size(segment.name), 4096); - } -} - -TEST(StoreResourceSnapshotCodecTest, - SnapshotRejectsUnsupportedDriversAndCorruptContainers) { - RegionDriverConfig config; - config.memory_allocator = BufferAllocatorType::CACHELIB; - auto drivers = CreateRegionDrivers(config); - ASSERT_TRUE(drivers.has_value()); - SegmentPool unsupported(std::move(*drivers)); - auto encoded = CaptureAndEncode(unsupported, {}); - ASSERT_FALSE(encoded.has_value()); - EXPECT_EQ(encoded.error().code, ErrorCode::SERIALIZE_UNSUPPORTED); - - SegmentPool cxl(Drivers(true)); - ASSERT_EQ(cxl.AcquireWriteAccess().MountSegment( - MakeSegment(0, "snapshot-cxl", "cxl"), generate_uuid()), - ErrorCode::OK); - encoded = CaptureAndEncode(cxl, {}); - ASSERT_FALSE(encoded.has_value()); - EXPECT_EQ(encoded.error().code, ErrorCode::SERIALIZE_UNSUPPORTED); - - SegmentPool pool(Drivers()); - const auto segment = MakeSegment(1, "snapshot-container"); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, generate_uuid()), - ErrorCode::OK); - auto allocator = pool.AcquireReadAccess().GetAllocator(segment.id); - - // Invalid zstd, invalid MessagePack, and a valid non-map root. - const uint8_t invalid_msgpack = 0xc1; - const uint8_t array = 0x90; - const std::vector> corrupt{ - {}, - {1, 2, 3}, - zstd_compress(&invalid_msgpack, 1, 3), - zstd_compress(&array, 1, 3)}; - for (const auto& bytes : corrupt) { - auto decoded = DecodeAndRestore(pool, bytes, false); - ASSERT_FALSE(decoded.has_value()); - EXPECT_EQ(decoded.error().code, ErrorCode::DESERIALIZE_FAIL); - EXPECT_EQ(pool.AcquireReadAccess().GetAllocator(segment.id), allocator); - } -} - -// Capacities from the payload drive node storage allocation, so they must be -// bounded and validated before the allocator is materialized. -TEST(StoreResourceSnapshotCodecTest, CorruptAllocatorCapacityIsRejected) { - SegmentPool pool(Drivers()); - const auto segment = MakeSegment(0, "snapshot-capacity"); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, generate_uuid()), - ErrorCode::OK); - auto encoded = CaptureAndEncode(pool, {}); - ASSERT_TRUE(encoded.has_value()); - - // Element 1 of the __Allocator payload is current_capacity, element 2 is - // max_capacity and element 9 is freeOffset. - const std::vector> corruptions{ - {2, uint64_t{offset_allocator::OffsetAllocatorSnapshot::kMaxNodes} + 1}, - {2, std::numeric_limits::max()}, - {1, std::numeric_limits::max()}, - {9, std::numeric_limits::max()}, - }; - for (const auto& [element, value] : corruptions) { - SCOPED_TRACE(element); - auto corrupted = RewriteSnapshot(*encoded, [&](auto& root, auto&) { - // region[7] = OffsetBufferAllocator, [5] = OffsetAllocator, [5] = - // __Allocator. - auto& region = SnapshotField(root, "ms").via.map.ptr[0].val; - auto& allocator = region.via.array.ptr[7]; - auto& offset_allocator = allocator.via.array.ptr[5]; - auto& layout = offset_allocator.via.array.ptr[5]; - layout.via.array.ptr[element] = msgpack::object(value); - }); - auto decoded = StoreResourceSnapshotCodec::Decode(corrupted); - ASSERT_FALSE(decoded.has_value()); - EXPECT_EQ(decoded.error().code, ErrorCode::DESERIALIZE_FAIL); - EXPECT_NE(decoded.error().message.find("capacity"), std::string::npos); - EXPECT_NE(pool.AcquireReadAccess().Catalog().Find(segment.id), nullptr); - } -} - -// A detached snapshot of a multi-gigabyte segment is pure metadata: encoding, -// decoding and restoring it must never mount or touch the described memory. -TEST(StoreResourceSnapshotCodecTest, - LargeSegmentSnapshotRoundTripsWithoutPhysicalMemory) { - constexpr uint64_t kLargeBase = 0x200000000ULL; - constexpr size_t kLargeCapacity = 64ULL * 1024 * 1024 * 1024; - constexpr size_t kAllocation = 4096; - auto& metrics = MasterMetricManager::instance(); - const auto baseline_used = metrics.get_allocated_mem_size(); - const auto baseline_capacity = metrics.get_total_mem_capacity(); - - // Two live metadata nodes describe the whole 64GiB region: one 4KiB - // allocation and the free remainder. Use the real host-memory node ceiling - // (64Mi), not a small test ceiling: the former 16Mi limit must fail this - // round trip. reserve() is lazy; nothing is mapped at kLargeBase. - auto source = offset_allocator::OffsetAllocator::create( - kLargeBase, kLargeCapacity, /*init_capacity=*/2, - offset_allocator::OffsetAllocatorSnapshot::kMaxNodes); - ASSERT_NE(source, nullptr); - auto allocation = source->allocate(kAllocation); - ASSERT_TRUE(allocation.has_value()); - - auto segment = MakeSegment(3, "detached-large"); - segment.base = kLargeBase; - segment.size = kLargeCapacity; - MountedRegion mounted; - mounted.segment = segment; - mounted.client_id = generate_uuid(); - mounted.status = SegmentStatus::OK; - - SegmentPoolSnapshot snapshot; - snapshot.active_names = {segment.name}; - snapshot.regions.push_back(RegionSnapshot{ - .mounted = mounted, - .allocator = - OffsetBufferAllocatorSnapshot{ - .segment_name = segment.name, - .base = kLargeBase, - .capacity = kLargeCapacity, - .used_bytes = kAllocation, - .transport_endpoint = segment.te_endpoint, - .allocation_state = source->CaptureSnapshot(), - }, - }); - auto valid = snapshot.regions.front().allocator.Validate(); - ASSERT_TRUE(valid.has_value()) << valid.error(); - - auto encoded = StoreResourceSnapshotCodec::Encode(snapshot, {}); - ASSERT_TRUE(encoded.has_value()) << encoded.error().message; - auto decoded = StoreResourceSnapshotCodec::Decode(*encoded); - ASSERT_TRUE(decoded.has_value()) << decoded.error().message; - ASSERT_EQ(decoded->segment_pool.regions.size(), 1U); - auto& state = decoded->segment_pool.regions.front().allocator; - EXPECT_EQ(state.capacity, kLargeCapacity); - EXPECT_EQ(state.used_bytes, kAllocation); - EXPECT_EQ(state.allocation_state.allocated_num, 1U); - - auto restored = OffsetBufferAllocator::Restore(std::move(state)); - ASSERT_TRUE(restored.has_value()); - auto allocator = std::move(*restored); - EXPECT_EQ(allocator->capacity(), kLargeCapacity); - EXPECT_EQ(allocator->size(), kAllocation); - EXPECT_EQ(allocator->base(), kLargeBase); - EXPECT_EQ(allocator->getSegmentName(), segment.name); - // The free remainder kept its size and unit multiplier across the round - // trip, which is only possible if the persisted nodes were restored. - EXPECT_EQ(allocator->getLargestFreeRegion(), - source->storageReport().largestFreeRegion); - - allocator.reset(); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline_used); - EXPECT_EQ(metrics.get_total_mem_capacity(), baseline_capacity); -} - -// A node payload whose zstd frame header declares far more than the expected -// node bytes must be rejected before the decompression buffer is allocated. -TEST(StoreResourceSnapshotCodecTest, - OversizedNodeFrameIsRejectedBeforeAllocation) { - SegmentPool pool(Drivers()); - const auto segment = MakeSegment(0, "snapshot-frame-header"); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, generate_uuid()), - ErrorCode::OK); - auto encoded = CaptureAndEncode(pool, {}); - ASSERT_TRUE(encoded.has_value()) << encoded.error().message; - - // Minimal zstd frame header: magic, a descriptor selecting the 8-byte - // Frame_Content_Size field in single-segment mode, then 1TiB. No content - // follows, so only the declared size can be examined. - constexpr uint64_t kDeclaredBytes = 1ULL << 40; - std::vector frame{0x28, 0xB5, 0x2F, 0xFD, 0xE0}; - for (uint32_t shift = 0; shift < 64; shift += 8) { - frame.push_back(static_cast(kDeclaredBytes >> shift)); - } - - // Element 7 of the __Allocator payload is the node data and element 8 the - // free-node data. - for (const uint32_t element : {7U, 8U}) { - SCOPED_TRACE(element); - auto corrupted = RewriteSnapshot(*encoded, [&](auto& root, auto& zone) { - // region[7] = OffsetBufferAllocator, [5] = OffsetAllocator, [5] = - // __Allocator. - auto& layout = SnapshotField(root, "ms") - .via.map.ptr[0] - .val.via.array.ptr[7] - .via.array.ptr[5] - .via.array.ptr[5]; - layout.via.array.ptr[element] = msgpack::object( - msgpack::type::raw_ref( - reinterpret_cast(frame.data()), frame.size()), - zone); - }); - auto decoded = StoreResourceSnapshotCodec::Decode(corrupted); - ASSERT_FALSE(decoded.has_value()); - EXPECT_EQ(decoded.error().code, ErrorCode::DESERIALIZE_FAIL); - EXPECT_NE(decoded.error().message.find("exceeds maximum allowed"), - std::string::npos) - << decoded.error().message; - } -} - -// The inner allocator counters and unit multiplier must describe the inner -// layout, and Decode reports the specific mismatch instead of a generic -// rejection. -TEST(StoreResourceSnapshotCodecTest, CorruptAllocatorInternalsAreRejected) { - SegmentPool pool(Drivers()); - const auto segment = MakeSegment(0, "snapshot-internals"); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, generate_uuid()), - ErrorCode::OK); - auto allocator = pool.AcquireReadAccess().GetAllocator(segment.id); - auto buffer = allocator->allocate(4096); - ASSERT_NE(buffer, nullptr); - auto encoded = CaptureAndEncode(pool, {}); - ASSERT_TRUE(encoded.has_value()) << encoded.error().message; - const auto baseline_used = - MasterMetricManager::instance().get_allocated_mem_size(); - - struct Corruption { - const char* name; - // OffsetAllocator element: 1 = multiplier_bits, 3 = allocated_size, - // 4 = allocated_num. - uint32_t element; - uint64_t value; - const char* reason; - // allocate() records the requested bytes in the outer usage too, so an - // altered allocated_size has to be mirrored to reach the inner check. - bool mirror_usage; - }; - const std::vector corruptions{ - {"allocated_size smaller than the used nodes", 3, 0, "allocated_size", - true}, - {"allocated_num mismatch", 4, 0, "allocated_num", false}, - {"multiplier_bits beyond the shift range", 1, 64, "multiplier_bits", - false}, - }; - for (const auto& corruption : corruptions) { - SCOPED_TRACE(corruption.name); - auto corrupted = RewriteSnapshot(*encoded, [&](auto& root, auto&) { - // region[7] = OffsetBufferAllocator, [5] = OffsetAllocator; element - // 3 of the OffsetBufferAllocator is the outer used_bytes. - auto& allocator_field = - SnapshotField(root, "ms").via.map.ptr[0].val.via.array.ptr[7]; - auto& state = allocator_field.via.array.ptr[5]; - state.via.array.ptr[corruption.element] = - msgpack::object(corruption.value); - if (corruption.mirror_usage) { - allocator_field.via.array.ptr[3] = - msgpack::object(corruption.value); - } - }); - auto decoded = StoreResourceSnapshotCodec::Decode(corrupted); - ASSERT_FALSE(decoded.has_value()); - EXPECT_EQ(decoded.error().code, ErrorCode::DESERIALIZE_FAIL); - EXPECT_NE(decoded.error().message.find(corruption.reason), - std::string::npos) - << decoded.error().message; - } - - // Decoding is pure: the published allocator and the metrics are unchanged. - EXPECT_EQ(pool.AcquireReadAccess().GetAllocator(segment.id), allocator); - EXPECT_EQ(MasterMetricManager::instance().get_allocated_mem_size(), - baseline_used); -} - -// The OffsetAllocator wire deserializer installs through the snapshot factory, -// so a rejected snapshot reports the specific reason instead of a generic -// rejection. -TEST(StoreResourceSnapshotCodecTest, - LegacyAllocatorDeserializeReportsDetailedReason) { - SegmentPool pool(Drivers()); - const auto segment = MakeSegment(0, "snapshot-legacy-reason"); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, generate_uuid()), - ErrorCode::OK); - auto allocator = pool.AcquireReadAccess().GetAllocator(segment.id); - auto buffer = allocator->allocate(4096); - ASSERT_NE(buffer, nullptr); - auto encoded = CaptureAndEncode(pool, {}); - ASSERT_TRUE(encoded.has_value()) << encoded.error().message; - - auto corrupted = RewriteSnapshot(*encoded, [](auto& root, auto&) { - // Element 4 of the OffsetAllocator array is allocated_num; it cannot - // exceed the single used node. - auto& state = SnapshotField(root, "ms") - .via.map.ptr[0] - .val.via.array.ptr[7] - .via.array.ptr[5]; - state.via.array.ptr[4] = msgpack::object(uint64_t{2}); - }); - const auto bytes = zstd_decompress(corrupted); - auto unpacked = msgpack::unpack(reinterpret_cast(bytes.data()), - bytes.size()); - auto root = unpacked.get(); - auto& offset_allocator_object = SnapshotField(root, "ms") - .via.map.ptr[0] - .val.via.array.ptr[7] - .via.array.ptr[5]; - auto decoded = Serializer::deserialize( - offset_allocator_object); - ASSERT_FALSE(decoded.has_value()); - EXPECT_EQ(decoded.error().code, ErrorCode::DESERIALIZE_FAIL); - EXPECT_NE(decoded.error().message.find("allocated_num"), std::string::npos) - << decoded.error().message; - - // Conversion failures must use the same expected-based error boundary, - // rather than leaking msgpack exceptions from the legacy entry point. - offset_allocator_object.via.array.ptr[4] = msgpack::object(); - decoded = Serializer::deserialize( - offset_allocator_object); - ASSERT_FALSE(decoded.has_value()); - EXPECT_EQ(decoded.error().code, ErrorCode::DESERIALIZE_FAIL); - EXPECT_NE(decoded.error().message.find("failed"), std::string::npos); - - auto& layout = offset_allocator_object.via.array.ptr[5]; - layout.via.array.ptr[9] = msgpack::object(); - auto decoded_layout = - Serializer::deserialize(layout); - ASSERT_FALSE(decoded_layout.has_value()); - EXPECT_EQ(decoded_layout.error().code, ErrorCode::DESERIALIZE_FAIL); - EXPECT_NE(decoded_layout.error().message.find("invalid scalar field"), - std::string::npos); -} - -} // namespace mooncake::ha diff --git a/mooncake-store/tests/local_ssd/local_ssd_test.cpp b/mooncake-store/tests/local_ssd/local_ssd_test.cpp index 6abe394c71..afb496df95 100644 --- a/mooncake-store/tests/local_ssd/local_ssd_test.cpp +++ b/mooncake-store/tests/local_ssd/local_ssd_test.cpp @@ -214,31 +214,5 @@ TEST(LocalSsdManagerTest, CancelsOffloadMirrorsOnlyWhenAllArePending) { EXPECT_TRUE(manager.SetOffloadingAndTakePending(client_b, true)->empty()); } -TEST(LocalSsdManagerTest, RestoreResetsRuntimeOnlyState) { - LocalSsdManager manager; - UUID client{5, 6}; - ASSERT_TRUE(manager.RegisterClient(client, true) == ErrorCode::OK); - ASSERT_TRUE(manager.ReportCapacity(client, 2048).has_value()); - ASSERT_TRUE(manager.AdjustUsedBytes(client, 512)); - ASSERT_TRUE(manager.EnqueueOffload(client, Offload("tenant", "key", 10), - 10) == ErrorCode::OK); - ASSERT_TRUE(manager.EnqueuePromotion( - client, Promotion("tenant", "key", 10)) == ErrorCode::OK); - manager.RequestRemoveAll(); - - auto state = manager.ExportPersistedState(); - manager.RestorePersistedState(std::move(state)); - auto usage = manager.GetUsage(client); - ASSERT_TRUE(usage.has_value()); - EXPECT_EQ(usage->total_capacity_bytes, 2048); - EXPECT_EQ(usage->used_bytes, 0); - EXPECT_TRUE(manager.TakePromotions(client, 10)->empty()); - EXPECT_FALSE(*manager.ConsumeRemoveAll(client)); - auto offloads = manager.SetOffloadingAndTakePending(client, true); - ASSERT_TRUE(offloads.has_value()); - ASSERT_EQ(offloads->size(), 1); - EXPECT_EQ(offloads->front(), Offload("tenant", "key", 10)); -} - } // namespace } // namespace mooncake diff --git a/mooncake-store/tests/master_metrics_test.cpp b/mooncake-store/tests/master_metrics_test.cpp index 6c25b5c593..b7670006bc 100644 --- a/mooncake-store/tests/master_metrics_test.cpp +++ b/mooncake-store/tests/master_metrics_test.cpp @@ -333,82 +333,6 @@ TEST_F(MasterMetricsTest, ServiceTeardownReleasesSegmentCapacity) { ASSERT_EQ(metrics.get_segment_total_mem_capacity(segment.name), 0); } -TEST_F(MasterMetricsTest, SnapshotReaderTeardownKeepsCapacityIntact) { - auto& metrics = MasterMetricManager::instance(); - - // Mount a segment through the accounted path so the gauge is non-zero. - SegmentManager source_manager(BufferAllocatorType::OFFSET); - Segment segment; - segment.id = generate_uuid(); - segment.name = "snapshot_reader_segment"; - segment.base = 0x300000000; - segment.size = 1024 * 1024 * 16; - UUID client_id = generate_uuid(); - ASSERT_EQ(source_manager.getSegmentAccess().MountSegment( - segment, client_id, - std::make_shared( - ClientLivenessRecord::Clock::now())), - ErrorCode::OK); - const int64_t capacity_after_mount = metrics.get_total_mem_capacity(); - ASSERT_EQ(metrics.get_segment_total_mem_capacity(segment.name), - static_cast(segment.size)); - - constexpr size_t kAllocationSize = 4 * 1024 * 1024; - std::shared_ptr source_allocator; - { - auto access = source_manager.getSegmentAccess(); - source_allocator = access.GetAllocator(segment.id); - } - ASSERT_NE(source_allocator, nullptr); - auto source_buffer = source_allocator->allocate(kAllocationSize); - ASSERT_NE(source_buffer, nullptr); - - auto snapshot = - SegmentSerializer(&source_manager).Serialize(LocalSsdPersistedState{}); - ASSERT_TRUE(snapshot.has_value()); - - { - // Deserialize into a temporary reader. The reader's records never - // contributed to the capacity metrics, so destroying it must leave - // the gauges untouched. - SegmentManager reader(BufferAllocatorType::OFFSET); - SegmentSerializer reader_serializer(&reader); - ASSERT_TRUE( - reader_serializer.Deserialize(snapshot.value()).has_value()); - const auto restored_usage = reader.GetMemoryUsageSnapshot(); - const auto restored_aggregate = reader.GetMemoryUsage(); - EXPECT_EQ(restored_usage.used_bytes, kAllocationSize); - EXPECT_EQ(restored_usage.capacity_bytes, segment.size); - EXPECT_DOUBLE_EQ(restored_usage.used_ratio(), 0.25); - EXPECT_EQ(restored_aggregate.used_bytes, kAllocationSize); - EXPECT_EQ(restored_aggregate.capacity_bytes, segment.size); - EXPECT_DOUBLE_EQ(restored_aggregate.used_ratio(), 0.25); - ASSERT_EQ(restored_usage.segments.size(), 1u); - EXPECT_EQ(restored_usage.segments.at(segment.name).used_bytes, - kAllocationSize); - - reader_serializer.Reset(); - EXPECT_EQ(reader.GetMemoryUsage().used_bytes, 0u); - EXPECT_EQ(reader.GetMemoryUsage().capacity_bytes, 0u); - } - ASSERT_EQ(metrics.get_total_mem_capacity(), capacity_after_mount); - ASSERT_EQ(metrics.get_segment_total_mem_capacity(segment.name), - static_cast(segment.size)); - - // Unmount to restore the gauges for the other tests. - source_buffer.reset(); - { - auto access = source_manager.getSegmentAccess(); - size_t dec_capacity = 0; - ASSERT_EQ(access.PrepareUnmountSegment(segment.id, dec_capacity), - ErrorCode::OK); - ASSERT_EQ( - access.CommitUnmountSegment(segment.id, client_id, dec_capacity), - ErrorCode::OK); - } - ASSERT_EQ(metrics.get_segment_total_mem_capacity(segment.name), 0); -} - TEST_F(MasterMetricsTest, CalcCacheStatsTest) { const uint64_t default_kv_lease_ttl = 100; auto& metrics = MasterMetricManager::instance(); diff --git a/mooncake-store/tests/offset_allocator_test.cpp b/mooncake-store/tests/offset_allocator_test.cpp index cb2273c102..abf976a2d5 100644 --- a/mooncake-store/tests/offset_allocator_test.cpp +++ b/mooncake-store/tests/offset_allocator_test.cpp @@ -1,4 +1,3 @@ -#include "allocator.h" #include "offset_allocator/offset_allocator.h" #include "mutex.h" #include "serializer.h" @@ -7,7 +6,6 @@ #include #include -#include #include #include #include @@ -344,171 +342,6 @@ class OffsetAllocatorTest : public ::testing::Test { handle.getHandle().m_allocator = new_allocator; } - // Corruptions used by SnapshotValidationRejectsInconsistentLayout. TEST_F - // bodies live in a derived class, so the layout access stays in fixture - // members, which hold __Allocator's friendship. - enum class LayoutFault { - kNone, - kBinIndexOutOfRange, - kUsedNodeInBinList, - kFreeStorageMismatch, - kFreeOffsetTooLarge, - kFreeOffsetTooSmall, - kDuplicateSpareNode, - kSpareNodeIsLive, - kUnclassifiedNode, - kNodeOffsetShifted, - kNeighborChainBroken, - kBinBitWithoutList, - kUsedNodeZeroSize, - kUsedNodeNotBinBoundary, - kBinCycle, - kNeighborOutOfRange, - kNodeCapacityZero, - kNodeCapacityTooLarge, - kFreeOffsetOutOfRange, - kNodeStorageTooShort, - kNodeStorageTooLong, - kFreeNodeStorageTooLong, - }; - - static uint32 firstUsedNode(const __Allocator& layout) { - for (uint32 i = 0; i < layout.m_current_capacity; ++i) { - if (layout.m_nodes[i].used) { - return i; - } - } - return __Allocator::Node::unused; - } - - static void applyLayoutFault(__Allocator& layout, LayoutFault fault) { - const uint32 used = firstUsedNode(layout); - ASSERT_NE(used, __Allocator::Node::unused); - switch (fault) { - case LayoutFault::kNone: - break; - case LayoutFault::kUsedNodeZeroSize: - layout.m_nodes[used].dataSize = 0; - break; - case LayoutFault::kUsedNodeNotBinBoundary: - ++layout.m_nodes[used].dataSize; - break; - case LayoutFault::kBinCycle: - for (uint32 bin = 0; bin < NUM_LEAF_BINS; ++bin) { - const uint32 head = layout.m_binIndices[bin]; - if (head != __Allocator::Node::unused) { - layout.m_nodes[head].binListNext = head; - return; - } - } - FAIL() << "no free bin to corrupt"; - break; - case LayoutFault::kNeighborOutOfRange: - layout.m_nodes[used].neighborNext = layout.m_current_capacity; - break; - case LayoutFault::kNodeCapacityZero: - layout.m_current_capacity = 0; - break; - case LayoutFault::kNodeCapacityTooLarge: - layout.m_max_capacity = OffsetAllocatorSnapshot::kMaxNodes + 1; - break; - case LayoutFault::kFreeOffsetOutOfRange: - layout.m_freeOffset = layout.m_current_capacity + 1; - break; - case LayoutFault::kNodeStorageTooShort: - layout.m_nodes.clear(); - break; - case LayoutFault::kNodeStorageTooLong: - layout.m_nodes.emplace_back(); - break; - case LayoutFault::kFreeNodeStorageTooLong: - layout.m_freeNodes.push_back(0); - break; - case LayoutFault::kBinIndexOutOfRange: - layout.m_binIndices[0] = layout.m_current_capacity; - break; - case LayoutFault::kUsedNodeInBinList: { - uint32 empty_bin = NUM_LEAF_BINS; - for (uint32 bin = 0; bin < NUM_LEAF_BINS; ++bin) { - if (layout.m_binIndices[bin] == __Allocator::Node::unused) { - empty_bin = bin; - break; - } - } - ASSERT_NE(empty_bin, NUM_LEAF_BINS); - layout.m_binIndices[empty_bin] = used; - layout.m_nodes[used].binListNext = __Allocator::Node::unused; - layout.m_nodes[used].binListPrev = __Allocator::Node::unused; - layout.m_usedBins[empty_bin >> TOP_BINS_INDEX_SHIFT] |= - uint8{1} << (empty_bin & LEAF_BINS_INDEX_MASK); - layout.m_usedBinsTop |= uint32{1} - << (empty_bin >> TOP_BINS_INDEX_SHIFT); - break; - } - case LayoutFault::kFreeStorageMismatch: - ++layout.m_freeStorage; - break; - case LayoutFault::kFreeOffsetTooLarge: - ++layout.m_freeOffset; - break; - case LayoutFault::kFreeOffsetTooSmall: - --layout.m_freeOffset; - break; - case LayoutFault::kDuplicateSpareNode: - ASSERT_LT(layout.m_freeOffset + 1, layout.m_current_capacity); - layout.m_freeNodes[layout.m_freeOffset + 1] = - layout.m_freeNodes[layout.m_freeOffset]; - break; - case LayoutFault::kSpareNodeIsLive: - layout.m_freeNodes[layout.m_freeOffset] = used; - break; - case LayoutFault::kUnclassifiedNode: - layout.m_nodes[used].used = false; - break; - case LayoutFault::kNodeOffsetShifted: - ++layout.m_nodes[used].dataOffset; - break; - case LayoutFault::kNeighborChainBroken: - layout.m_nodes[used].neighborNext = __Allocator::Node::unused; - break; - case LayoutFault::kBinBitWithoutList: - for (uint32 bin = 0; bin < NUM_LEAF_BINS; ++bin) { - if (layout.m_binIndices[bin] == __Allocator::Node::unused) { - layout.m_usedBins[bin >> TOP_BINS_INDEX_SHIFT] |= - uint8{1} << (bin & LEAF_BINS_INDEX_MASK); - return; - } - } - FAIL() << "no empty bin to corrupt"; - break; - } - } - - static bool layoutPassesValidation( - const std::shared_ptr& allocator, LayoutFault fault) { - auto snapshot = allocator->CaptureSnapshot(); - applyLayoutFault(*snapshot.layout, fault); - return snapshot.Validate().has_value(); - } - - // End-to-end: the store level factory must refuse to install the same - // broken layout the validator rejects. - static ErrorCode restoreBufferAllocatorAfterLayoutFault( - const std::shared_ptr& allocator, LayoutFault fault) { - auto state = allocator->CaptureSnapshot(); - applyLayoutFault(*state.layout, fault); - OffsetBufferAllocatorSnapshot snapshot{ - .segment_name = "segment", - .base = static_cast(state.base), - .capacity = state.capacity, - .used_bytes = state.allocated_size, - .transport_endpoint = "endpoint", - .allocation_state = std::move(state), - }; - auto restored = OffsetBufferAllocator::Restore(std::move(snapshot)); - return restored ? ErrorCode::OK : restored.error(); - } - void assertAllocatorEQ(const std::shared_ptr& a, const std::shared_ptr& b) { MutexLocker lock_a(&a->m_mutex); @@ -2024,173 +1857,6 @@ TEST_F(OffsetAllocatorTest, ChainedAllocationAndDeserialization) { } } -// A persisted layout is only installed when it is internally consistent, so a -// corrupt layout cannot produce overlapping allocations, dangling bin links or -// a free stack that hands the same node out twice. Both snapshot validation -// and the lowest-level restore factory reject the same corruptions. -TEST_F(OffsetAllocatorTest, SnapshotValidationRejectsInconsistentLayout) { - constexpr uint64_t kBase = 0x100000000ULL; - constexpr size_t kSize = 16U * 1024 * 1024; - auto allocator = OffsetAllocator::create(kBase, kSize); - ASSERT_NE(allocator, nullptr); - auto first = allocator->allocate(4096); - auto second = allocator->allocate(4096); - ASSERT_TRUE(first.has_value()); - ASSERT_TRUE(second.has_value()); - - auto clean = allocator->CaptureSnapshot(); - ASSERT_TRUE(clean.Validate().has_value()); - EXPECT_EQ(clean.allocated_num, 2U); - EXPECT_EQ(clean.allocated_size, 8192U); - - const std::vector faults{ - LayoutFault::kBinIndexOutOfRange, - LayoutFault::kUsedNodeInBinList, - LayoutFault::kFreeStorageMismatch, - LayoutFault::kFreeOffsetTooLarge, - LayoutFault::kFreeOffsetTooSmall, - LayoutFault::kDuplicateSpareNode, - LayoutFault::kSpareNodeIsLive, - LayoutFault::kUnclassifiedNode, - LayoutFault::kNodeOffsetShifted, - LayoutFault::kNeighborChainBroken, - LayoutFault::kBinBitWithoutList, - LayoutFault::kUsedNodeZeroSize, - LayoutFault::kUsedNodeNotBinBoundary, - LayoutFault::kBinCycle, - LayoutFault::kNeighborOutOfRange, - LayoutFault::kNodeCapacityZero, - LayoutFault::kNodeCapacityTooLarge, - LayoutFault::kFreeOffsetOutOfRange, - LayoutFault::kNodeStorageTooShort, - LayoutFault::kNodeStorageTooLong, - LayoutFault::kFreeNodeStorageTooLong, - }; - for (const auto fault : faults) { - SCOPED_TRACE(static_cast(fault)); - EXPECT_FALSE(layoutPassesValidation(allocator, fault)); - auto snapshot = allocator->CaptureSnapshot(); - applyLayoutFault(*snapshot.layout, fault); - const auto valid = snapshot.Validate(); - ASSERT_FALSE(valid.has_value()); - EXPECT_FALSE(valid.error().empty()); - auto restored = OffsetAllocator::Restore(std::move(snapshot)); - ASSERT_FALSE(restored.has_value()); - EXPECT_EQ(restored.error(), valid.error()); - EXPECT_EQ(restoreBufferAllocatorAfterLayoutFault(allocator, fault), - ErrorCode::INVALID_PARAMS); - } - - // Rejected layouts must not disturb the allocator they came from. - EXPECT_TRUE(layoutPassesValidation(allocator, LayoutFault::kNone)); -} - -TEST_F(OffsetAllocatorTest, - SnapshotValidationRejectsHeaderAndAccountingFaults) { - auto allocator = OffsetAllocator::create(0x100000000ULL, 16U << 20, 16, 32); - auto first = allocator->allocate(4096); - auto second = allocator->allocate(4096); - ASSERT_TRUE(first.has_value()); - ASSERT_TRUE(second.has_value()); - struct Fault { - const char* diagnostic; - std::function apply; - }; - const std::vector faults{ - {"layout", [](auto& s) { s.layout.reset(); }}, - {"bounds", [](auto& s) { s.capacity = 0; }}, - {"bounds", - [](auto& s) { s.base = std::numeric_limits::max(); }}, - {"multiplier_bits", [](auto& s) { s.multiplier_bits = 64; }}, - {"multiplier_bits", - [](auto& s) { - s.multiplier_bits = std::numeric_limits::max(); - }}, - {"multiplier_bits", [](auto& s) { s.multiplier_bits = 1; }}, - {"layout size", [](auto& s) { ++s.capacity; }}, - {"allocated_num", [](auto& s) { s.allocated_num = 0; }}, - {"allocated_num", [](auto& s) { ++s.allocated_num; }}, - {"allocated_size", [](auto& s) { s.allocated_size = 0; }}, - {"allocated_size", [](auto& s) { s.allocated_size = 7681; }}, - {"allocated_size", [](auto& s) { s.allocated_size = 8193; }}, - }; - for (const auto& fault : faults) { - SCOPED_TRACE(fault.diagnostic); - auto snapshot = allocator->CaptureSnapshot(); - fault.apply(snapshot); - const auto valid = snapshot.Validate(); - ASSERT_FALSE(valid.has_value()); - EXPECT_NE(valid.error().find(fault.diagnostic), std::string::npos); - auto restored = OffsetAllocator::Restore(std::move(snapshot)); - ASSERT_FALSE(restored.has_value()); - EXPECT_EQ(restored.error(), valid.error()); - } -} - -TEST_F(OffsetAllocatorTest, SnapshotCapacityPreflightRejectsInvalidFields) { - const auto validate = OffsetAllocatorSnapshot::ValidateNodeCapacity; - EXPECT_TRUE(validate(1024, 16, OffsetAllocatorSnapshot::kMaxNodes, 1)); - EXPECT_TRUE(validate(1024, OffsetAllocatorSnapshot::kMaxNodes, - OffsetAllocatorSnapshot::kMaxNodes, - OffsetAllocatorSnapshot::kMaxNodes)); - EXPECT_FALSE(validate(0, 16, 32, 1)); - EXPECT_FALSE(validate(1024, 0, 32, 0)); - EXPECT_FALSE(validate(1024, 16, 0, 1)); - EXPECT_FALSE(validate(1024, 33, 32, 1)); - EXPECT_FALSE(validate(1024, 16, OffsetAllocatorSnapshot::kMaxNodes + 1, 1)); - EXPECT_FALSE(validate(1024, 16, 32, 17)); -} - -TEST_F(OffsetAllocatorTest, SnapshotValidationAcceptsRoundingAndBoundsUsage) { - for (const uint64_t capacity : {uint64_t{16} << 20, uint64_t{32} << 30}) { - SCOPED_TRACE(capacity); - auto allocator = - OffsetAllocator::create(0x100000000ULL, capacity, 16, 32); - auto allocation = allocator->allocate(4097); - ASSERT_TRUE(allocation.has_value()); - auto snapshot = allocator->CaptureSnapshot(); - ASSERT_TRUE(snapshot.Validate().has_value()); - EXPECT_EQ(snapshot.allocated_size, 4097U); - auto restored = OffsetAllocator::Restore(std::move(snapshot)); - ASSERT_TRUE(restored.has_value()); - EXPECT_EQ((*restored)->get_metrics().allocated_size_, 4097U); - EXPECT_EQ((*restored)->get_metrics().allocated_num_, 1U); - -#ifndef OFFSET_ALLOCATOR_NOT_ROUND_UP - // A 4097-byte request consumes 4608 bytes. Both bin rounding and the - // large allocator's unit shift allow [4097, 4608], not arbitrary values - // below the occupied size. Exact original requests cannot be recovered - // from the layout alone. - for (const uint64_t requested : {4096U, 4097U, 4608U, 4609U}) { - auto candidate = allocator->CaptureSnapshot(); - candidate.allocated_size = requested; - EXPECT_EQ(candidate.Validate().has_value(), - requested >= 4097 && requested <= 4608); - } -#endif - } -} - -TEST_F(OffsetAllocatorTest, - SnapshotValidationAcceptsEmptyFullAndSmallAllocations) { - auto allocator = OffsetAllocator::create(0, 16, 16, 32); - EXPECT_TRUE(allocator->CaptureSnapshot().Validate().has_value()); - auto full = allocator->allocate(16); - ASSERT_TRUE(full.has_value()); - EXPECT_TRUE(allocator->CaptureSnapshot().Validate().has_value()); - full.reset(); - EXPECT_TRUE(allocator->CaptureSnapshot().Validate().has_value()); - - for (const uint64_t capacity : {uint64_t{16} << 20, uint64_t{32} << 30}) { - auto small = OffsetAllocator::create(0, capacity, 16, 32); - for (const size_t requested : {1U, 7U, 8U, 9U}) { - auto allocation = small->allocate(requested); - ASSERT_TRUE(allocation.has_value()); - EXPECT_TRUE(small->CaptureSnapshot().Validate().has_value()); - } - } -} - } // namespace mooncake::offset_allocator int main(int argc, char** argv) { diff --git a/mooncake-store/tests/segment_pool_test.cpp b/mooncake-store/tests/segment_pool_test.cpp index eda1f50a2e..c27afb0c51 100644 --- a/mooncake-store/tests/segment_pool_test.cpp +++ b/mooncake-store/tests/segment_pool_test.cpp @@ -1,13 +1,8 @@ #include "segment/pool.h" #include -#include - -#include -#include #include -#include #include #include #include @@ -18,9 +13,7 @@ #include #include "client_liveness.h" -#include "ha/snapshot/store_resource_snapshot_codec.h" #include "master_metric_manager.h" -#include "segment.h" #include "segment/pool_read_access.h" #include "segment/pool_write_access.h" #include "test_buffer_allocator.h" @@ -79,29 +72,6 @@ void CommitUnmount(SegmentPool& pool, const Segment& segment, ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); } -tl::expected, SerializationError> CaptureAndEncode( - const SegmentPool& pool, const LocalSsdPersistedState& local_ssd) { - auto snapshot = pool.CaptureSnapshot(); - if (!snapshot) { - return tl::make_unexpected( - SerializationError(snapshot.error(), "capture SegmentPool failed")); - } - return ha::StoreResourceSnapshotCodec::Encode(*snapshot, local_ssd); -} - -tl::expected DecodeAndRestore( - SegmentPool& pool, const std::vector& bytes, - bool account_capacity) { - auto decoded = ha::StoreResourceSnapshotCodec::Decode(bytes); - if (!decoded) return tl::make_unexpected(decoded.error()); - auto restored = pool.RestoreSnapshot(std::move(decoded->segment_pool), - account_capacity); - if (!restored) - return tl::make_unexpected( - SerializationError(restored.error(), "restore SegmentPool failed")); - return std::move(decoded->local_ssd); -} - } // namespace TEST(SegmentPoolTest, ClassifiesOnlyCxlProtocolAsCxl) { @@ -1042,483 +1012,44 @@ TEST(SegmentPoolTest, AdoptRejectsCatalogConflictWithoutReplacingResource) { EXPECT_EQ(access.Catalog().Regions().size(), 1U); } -TEST(SegmentPoolTest, SnapshotCaptureDoesNotAcquireInheritedPoolMutex) { - SegmentPool pool(Drivers()); - const UUID client = generate_uuid(); - auto segment = MakeSegment(0, "fork-safe"); - { - auto access = pool.AcquireWriteAccess(); - ASSERT_EQ(access.MountSegment(segment, client), ErrorCode::OK); - - // A fork child can inherit this mutex as locked by a vanished thread. - // Capture must not acquire the inherited runtime lock. - auto encoded = CaptureAndEncode(pool, LocalSsdPersistedState{}); - ASSERT_TRUE(encoded.has_value()) << encoded.error().message; - } - CommitUnmount(pool, segment, client); -} - -TEST(SegmentPoolTest, SnapshotRoundTripPreservesCatalogAndHost) { - SegmentPool source(Drivers()); - const UUID client = generate_uuid(); - auto segment = MakeSegment(0, "snapshot", "tcp", "host-a"); - ASSERT_EQ(source.AcquireWriteAccess().MountSegment(segment, client), - ErrorCode::OK); - auto encoded = CaptureAndEncode(source, LocalSsdPersistedState{}); - ASSERT_TRUE(encoded.has_value()) << encoded.error().message; - - SegmentPool restored(Drivers()); - auto stale = MakeSegment(1, "stale"); - ASSERT_EQ(restored.AcquireWriteAccess().MountSegment(stale, client), - ErrorCode::OK); - // Replacement must also allow an existing UUID and a different old owner. - ASSERT_EQ( - restored.AcquireWriteAccess().MountSegment(segment, generate_uuid()), - ErrorCode::OK); - auto decoded = DecodeAndRestore(restored, *encoded, false); - ASSERT_TRUE(decoded.has_value()) << decoded.error().message; - - { - auto view = restored.AcquireReadAccess(); - EXPECT_EQ(view.Catalog().Find(stale.id), nullptr); - const auto* mounted = view.Catalog().Find(segment.id); - ASSERT_NE(mounted, nullptr); - EXPECT_EQ(mounted->segment.host_id, segment.host_id); - EXPECT_EQ(mounted->client_id, client); - EXPECT_EQ(mounted->status, SegmentStatus::OK); - } - - restored.AcquireWriteAccess().Clear(); - CommitUnmount(source, segment, client); -} - -TEST(SegmentPoolTest, RestoreSnapshotRejectsConflictsBeforeReplacingPool) { - SegmentPool pool(Drivers()); - const auto segment = MakeSegment(0, "restore-conflict"); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, generate_uuid()), - ErrorCode::OK); - auto allocator = pool.AcquireReadAccess().GetAllocator(segment.id); - auto snapshot = pool.CaptureSnapshot(); - auto duplicate = pool.CaptureSnapshot(); - ASSERT_TRUE(snapshot); - ASSERT_TRUE(duplicate); - snapshot->regions.push_back(std::move(duplicate->regions.front())); - EXPECT_FALSE(pool.RestoreSnapshot(std::move(*snapshot), true)); - EXPECT_EQ(pool.AcquireReadAccess().GetAllocator(segment.id), allocator); - - auto invalid = pool.CaptureSnapshot(); - ASSERT_TRUE(invalid); - invalid->regions.front().allocator.allocation_state.layout.reset(); - EXPECT_FALSE(pool.RestoreSnapshot(std::move(*invalid), true)); - EXPECT_EQ(pool.AcquireReadAccess().GetAllocator(segment.id), allocator); -} - -TEST(SegmentPoolTest, SnapshotReplacementPreservesAllocatedMetrics) { - auto& metrics = MasterMetricManager::instance(); - const auto baseline = metrics.get_allocated_mem_size(); - - for (bool account_capacity : {false, true}) { - SegmentPool pool(Drivers()); - const UUID client = generate_uuid(); - auto segment = MakeSegment(0, "snapshot-allocated-metrics"); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, client), - ErrorCode::OK); - auto buffer = - pool.AcquireReadAccess().GetAllocator(segment.id)->allocate(4096); - ASSERT_NE(buffer, nullptr); - auto encoded = CaptureAndEncode(pool, LocalSsdPersistedState{}); - ASSERT_TRUE(encoded.has_value()) << encoded.error().message; - - auto decoded = DecodeAndRestore(pool, *encoded, account_capacity); - ASSERT_TRUE(decoded.has_value()) << decoded.error().message; - buffer.reset(); // The replaced allocator is no longer live. - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline + 4096); - EXPECT_EQ(metrics.get_segment_allocated_mem_size(segment.name), 4096); - EXPECT_EQ(metrics.get_segment_total_mem_capacity(segment.name), - account_capacity ? kRegionSize : 0); - { - auto next = pool.AcquireReadAccess() - .GetAllocator(segment.id) - ->allocate(4096); - ASSERT_NE(next, nullptr); - EXPECT_EQ(metrics.get_segment_allocated_mem_size(segment.name), - 8192); - } - EXPECT_EQ(metrics.get_segment_allocated_mem_size(segment.name), 4096); - pool.AcquireWriteAccess().Clear(); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline); - EXPECT_EQ(metrics.get_segment_allocated_mem_size(segment.name), 0); - } -} - -TEST(SegmentPoolTest, SnapshotRestoresPlacementAndEveryLifecycleStatus) { - SegmentPool source(Drivers()); - const UUID client = generate_uuid(); - std::vector segments; - const std::vector statuses{ - SegmentStatus::OK, SegmentStatus::OK, - SegmentStatus::UNDEFINED, SegmentStatus::DRAINING, - SegmentStatus::DRAINED, SegmentStatus::GRACEFULLY_UNMOUNTING, - SegmentStatus::UNMOUNTING}; - for (size_t i = 0; i < statuses.size(); ++i) { - segments.push_back(MakeSegment(i, "lifecycle-" + std::to_string(i), - "tcp", "host-" + std::to_string(i))); - auto access = source.AcquireWriteAccess(); - ASSERT_EQ( - access.MountSegment(segments.back(), i == 1 ? UUID{8, 9} : client), - ErrorCode::OK); - ASSERT_EQ( - access.SetSegmentStatusByName(segments.back().name, statuses[i]), - ErrorCode::OK); - } - - // Same-name regions must all survive; only their logical name is - // deduplicated. - auto sibling = MakeSegment(8, segments[0].name, "tcp", segments[0].host_id); - ASSERT_EQ(source.AcquireWriteAccess().MountSegment(sibling, client), - ErrorCode::OK); - auto encoded = CaptureAndEncode(source, {}); - ASSERT_TRUE(encoded.has_value()); - SegmentPool restored(Drivers()); - ASSERT_TRUE(DecodeAndRestore(restored, *encoded, false).has_value()); - { - auto access = restored.AcquireReadAccess(); - EXPECT_EQ(access.Catalog().Regions().size(), segments.size() + 1); - for (size_t i = 0; i < segments.size(); ++i) { - const auto* region = access.Catalog().Find(segments[i].id); - ASSERT_NE(region, nullptr); - EXPECT_EQ(region->status, statuses[i]); - EXPECT_EQ(region->client_id, i == 1 ? (UUID{8, 9}) : client); - EXPECT_EQ(access.Placement().Contains( - segments[i].name, AllocationCandidateKind::NATIVE), - statuses[i] == SegmentStatus::OK); - EXPECT_NE(access.GetAllocator(segments[i].id), nullptr); - } - std::vector names; - access.Placement().GetActiveSegmentNames( - AllocationCandidateKind::NATIVE, names); - EXPECT_EQ(names, (std::vector{segments[0].name, - segments[1].name})); - std::vector by_host; - access.Placement().VisitHostOrderedSegmentNames( - segments[1].host_id, "key", [&](auto name) { - by_host.emplace_back(name); - return false; - }); - ASSERT_EQ(by_host.size(), 2U); - EXPECT_EQ(by_host.front(), segments[1].name); - const auto* entry = access.Placement().Find( - segments[0].name, AllocationCandidateKind::NATIVE); - ASSERT_NE(entry, nullptr); - EXPECT_EQ(entry->candidates.size(), 2U); - } - - CommitUnmount(restored, segments[0], client); - EXPECT_TRUE(restored.AcquireReadAccess().Placement().Contains( - sibling.name, AllocationCandidateKind::NATIVE)); - CommitUnmount(restored, sibling, client); - EXPECT_FALSE(restored.AcquireReadAccess().Placement().Contains( - sibling.name, AllocationCandidateKind::NATIVE)); - auto buffer = restored.AcquireReadAccess() - .GetAllocator(segments[1].id) - ->allocate(4096); - ASSERT_NE(buffer, nullptr); -} - -TEST(SegmentPoolTest, - SnapshotPreservesFragmentationAndLegacyAllocationHandles) { - SegmentManager legacy(BufferAllocatorType::OFFSET); - const UUID client = generate_uuid(); - const auto segment = MakeSegment(0, "snapshot-handles"); - const auto client_liveness = std::make_shared( - ClientLivenessRecord::Clock::now()); - ASSERT_EQ(legacy.getSegmentAccess().MountSegment(segment, client, - client_liveness), - ErrorCode::OK); - - // SegmentManager does not release capacity in its destructor. - legacy.releaseCapacityMetrics(); - auto allocator = legacy.getSegmentAccess().GetAllocator(segment.id); - auto first = allocator->allocate(4096); - auto hole = allocator->allocate(4096); - auto last = allocator->allocate(4096); - ASSERT_NE(first, nullptr); - ASSERT_NE(hole, nullptr); - ASSERT_NE(last, nullptr); - hole.reset(); - - msgpack::sbuffer handles; - MsgpackPacker packer(&handles); - packer.pack_array(2); - ASSERT_TRUE( - Serializer::serialize(*first, legacy.getView(), packer) - .has_value()); - ASSERT_TRUE( - Serializer::serialize(*last, legacy.getView(), packer) - .has_value()); - - auto encoded = SegmentSerializer(&legacy).Serialize({}); - ASSERT_TRUE(encoded.has_value()); - SegmentPool restored(Drivers()); - ASSERT_TRUE(DecodeAndRestore(restored, *encoded, false).has_value()); - auto target = std::dynamic_pointer_cast( - restored.AcquireReadAccess().GetAllocator(segment.id)); - ASSERT_NE(target, nullptr); - EXPECT_EQ(restored.GetMemoryUsage().used_bytes, 8192U); - - auto next = target->allocate(4096); - ASSERT_NE(next, nullptr); - EXPECT_NE(next->data(), first->data()); - EXPECT_NE(next->data(), last->data()); - next.reset(); - - auto objects = msgpack::unpack(handles.data(), handles.size()); - for (uint32_t i = 0; i < 2; ++i) { - const auto& fields = objects.get().via.array.ptr[i].via.array; - auto handle = - Serializer::deserialize( - fields.ptr[4], target->getOffsetAllocator()); - ASSERT_TRUE(handle.has_value()); - auto buffer = std::make_unique( - target, reinterpret_cast(fields.ptr[1].as()), - fields.ptr[0].as(), std::move(**handle)); - buffer.reset(); - EXPECT_EQ(restored.GetMemoryUsage().used_bytes, (1U - i) * 4096U); - } - - // Restored handles released and coalesced every block. - auto all = target->allocate(kRegionSize); - ASSERT_NE(all, nullptr); -} - -TEST(SegmentPoolTest, SnapshotReadersAndEmptyReplacementBalanceMetrics) { - auto& metrics = MasterMetricManager::instance(); - SegmentPool source(Drivers()); - const auto segment = MakeSegment(0, "snapshot-reader"); - ASSERT_EQ( - source.AcquireWriteAccess().MountSegment(segment, generate_uuid()), - ErrorCode::OK); - auto live = - source.AcquireReadAccess().GetAllocator(segment.id)->allocate(4096); - ASSERT_NE(live, nullptr); - auto encoded = CaptureAndEncode(source, {}); - ASSERT_TRUE(encoded.has_value()); - - SegmentPool empty(Drivers()); - auto empty_encoded = CaptureAndEncode(empty, {}); - ASSERT_TRUE(empty_encoded.has_value()); - const auto baseline_used = metrics.get_allocated_mem_size(); - const auto baseline_capacity = metrics.get_total_mem_capacity(); - - for (bool account_capacity : {false, true}) { - for (bool clear_with_empty : {false, true}) { - { - SegmentPool reader(Drivers()); - for (int i = 0; i < 3; ++i) { - ASSERT_TRUE( - DecodeAndRestore(reader, *encoded, account_capacity) - .has_value()); - EXPECT_EQ(reader.GetMemoryUsage().used_bytes, 4096U); - EXPECT_EQ(reader.GetMemoryUsage().capacity_bytes, - kRegionSize); - EXPECT_EQ(metrics.get_allocated_mem_size(), - baseline_used + 4096); - EXPECT_EQ(metrics.get_total_mem_capacity(), - baseline_capacity + - (account_capacity ? kRegionSize : 0)); - } - - if (clear_with_empty) { - ASSERT_TRUE(DecodeAndRestore(reader, *empty_encoded, - account_capacity) - .has_value()); - EXPECT_TRUE( - reader.AcquireReadAccess().Catalog().Regions().empty()); - EXPECT_EQ(reader.GetMemoryUsage().used_bytes, 0U); - EXPECT_EQ(reader.GetMemoryUsage().capacity_bytes, 0U); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline_used); - } - } - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline_used); - EXPECT_EQ(metrics.get_total_mem_capacity(), baseline_capacity); - EXPECT_EQ(metrics.get_segment_allocated_mem_size(segment.name), - 4096); - EXPECT_EQ(metrics.get_segment_total_mem_capacity(segment.name), - kRegionSize); - } - } -} - -TEST(SegmentPoolTest, SnapshotReplacementBalancesExternallyRetainedAllocator) { - auto& metrics = MasterMetricManager::instance(); - const auto baseline = metrics.get_allocated_mem_size(); - SegmentPool pool(Drivers()); - const auto segment = MakeSegment(0, "snapshot-retained"); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, generate_uuid()), - ErrorCode::OK); - auto retained = pool.AcquireReadAccess().GetAllocator(segment.id); - auto old_buffer = retained->allocate(4096); - ASSERT_NE(old_buffer, nullptr); - auto encoded = CaptureAndEncode(pool, {}); - ASSERT_TRUE(encoded.has_value()); - ASSERT_TRUE(DecodeAndRestore(pool, *encoded, true).has_value()); - EXPECT_NE(pool.AcquireReadAccess().GetAllocator(segment.id), retained); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline + 8192); - EXPECT_EQ(metrics.get_segment_allocated_mem_size(segment.name), 8192); - - // Clearing the replacement must not erase the old allocator's usage label. - pool.AcquireWriteAccess().Clear(); - EXPECT_EQ(metrics.get_segment_total_mem_capacity(segment.name), 0); - EXPECT_EQ(metrics.get_segment_allocated_mem_size(segment.name), 4096); - - old_buffer.reset(); - retained.reset(); - EXPECT_EQ(metrics.get_segment_allocated_mem_size(segment.name), 0); - EXPECT_EQ(metrics.get_allocated_mem_size(), baseline); - EXPECT_EQ(pool.GetMemoryUsage().used_bytes, 0U); -} - TEST(SegmentPoolTest, DestroyedPoolAndRetainedAllocatorRemoveMetricLabels) { auto& metrics = MasterMetricManager::instance(); for (bool retain_allocator : {false, true}) { SCOPED_TRACE(retain_allocator); - const std::string name = "snapshot-label-cleanup"; + const std::string name = "retained-allocator-labels"; const std::string label = "segment=\"" + name + "\""; std::shared_ptr retained; + std::unique_ptr live; { - SegmentPool source(Drivers()); + SegmentPool pool(Drivers()); auto segment = MakeSegment(0, name); - ASSERT_EQ(source.AcquireWriteAccess().MountSegment(segment, - generate_uuid()), + ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, + generate_uuid()), ErrorCode::OK); - auto live = source.AcquireReadAccess() - .GetAllocator(segment.id) - ->allocate(4096); + auto allocator = pool.AcquireReadAccess().GetAllocator(segment.id); + ASSERT_NE(allocator, nullptr); + live = allocator->allocate(4096); ASSERT_NE(live, nullptr); - auto encoded = CaptureAndEncode(source, {}); - ASSERT_TRUE(encoded.has_value()); - SegmentPool restored(Drivers()); - ASSERT_TRUE(DecodeAndRestore(restored, *encoded, true).has_value()); if (retain_allocator) { - retained = - restored.AcquireReadAccess().GetAllocator(segment.id); + retained = std::move(allocator); } EXPECT_NE(metrics.serialize_metrics().find(label), std::string::npos); } + // The pool released its capacity, but a live allocator still owns + // usage under the label, so the label must outlive the pool. + EXPECT_EQ(metrics.get_segment_total_mem_capacity(name), 0); if (retain_allocator) { EXPECT_EQ(metrics.get_segment_allocated_mem_size(name), 4096); EXPECT_NE(metrics.serialize_metrics().find(label), std::string::npos); } retained.reset(); + live.reset(); EXPECT_EQ(metrics.get_segment_allocated_mem_size(name), 0); EXPECT_EQ(metrics.get_segment_total_mem_capacity(name), 0); EXPECT_EQ(metrics.serialize_metrics().find(label), std::string::npos); } } -TEST(SegmentPoolTest, EmptySnapshotAllocatorKeepsOtherPoolsCapacityLabels) { - SegmentPool source(Drivers()); - const auto segment = MakeSegment(0, "snapshot-shared-capacity"); - ASSERT_EQ( - source.AcquireWriteAccess().MountSegment(segment, generate_uuid()), - ErrorCode::OK); - auto encoded = CaptureAndEncode(source, {}); - ASSERT_TRUE(encoded.has_value()); - auto& metrics = MasterMetricManager::instance(); - const auto baseline = metrics.get_total_mem_capacity(); - { - SegmentPool reader(Drivers()); - ASSERT_TRUE(DecodeAndRestore(reader, *encoded, true).has_value()); - EXPECT_EQ(metrics.get_segment_total_mem_capacity(segment.name), - 2 * kRegionSize); - reader.AcquireWriteAccess().Clear(); - } - EXPECT_EQ(metrics.get_total_mem_capacity(), baseline); - EXPECT_EQ(metrics.get_segment_total_mem_capacity(segment.name), - kRegionSize); - source.AcquireWriteAccess().Clear(); - EXPECT_EQ(metrics.get_segment_total_mem_capacity(segment.name), 0); -} - -TEST(SegmentPoolTest, ForkedSnapshotEncodesWhileAnotherThreadOwnsPoolLock) { - SegmentPool pool(Drivers()); - const UUID client = generate_uuid(); - const auto segment = MakeSegment(0, "real-fork"); - ASSERT_EQ(pool.AcquireWriteAccess().MountSegment(segment, client), - ErrorCode::OK); - auto live = - pool.AcquireReadAccess().GetAllocator(segment.id)->allocate(4096); - ASSERT_NE(live, nullptr); - - int fds[2]; - ASSERT_EQ(pipe(fds), 0); - std::promise locked, release; - auto released = release.get_future(); - std::thread holder([&] { - auto access = pool.AcquireWriteAccess(); - locked.set_value(); - released.wait(); - }); - locked.get_future().wait(); - - const pid_t child = fork(); - if (child == 0) { - close(fds[0]); - alarm(10); // A lock regression must fail, not hang the test suite. - auto encoded = CaptureAndEncode(pool, {}); - if (!encoded) { - _exit(1); - } - - size_t offset = 0; - while (offset < encoded->size()) { - const auto n = write(fds[1], encoded->data() + offset, - encoded->size() - offset); - if (n < 0 && errno == EINTR) { - continue; - } - if (n <= 0) { - _exit(2); - } - offset += static_cast(n); - } - _exit(0); - } - - close(fds[1]); - release.set_value(); - holder.join(); - - std::vector bytes; - uint8_t chunk[4096]; - ssize_t n; - while ((n = read(fds[0], chunk, sizeof(chunk))) != 0) { - if (n < 0 && errno == EINTR) { - continue; - } - if (n < 0) { - break; - } - bytes.insert(bytes.end(), chunk, chunk + n); - } - close(fds[0]); - ASSERT_GT(child, 0); - - int status = 0; - pid_t waited; - do { - waited = waitpid(child, &status, 0); - } while (waited < 0 && errno == EINTR); - ASSERT_EQ(waited, child); - ASSERT_TRUE(WIFEXITED(status)) << status; - ASSERT_EQ(WEXITSTATUS(status), 0); - - SegmentPool restored(Drivers()); - ASSERT_TRUE(DecodeAndRestore(restored, bytes, false).has_value()); - EXPECT_EQ(restored.GetMemoryUsage().used_bytes, 4096U); - EXPECT_NE(restored.AcquireReadAccess().Catalog().Find(segment.id), nullptr); -} - } // namespace mooncake::test diff --git a/mooncake-store/tests/segment_test.cpp b/mooncake-store/tests/segment_test.cpp index 942acf7fba..9c2170e5a1 100644 --- a/mooncake-store/tests/segment_test.cpp +++ b/mooncake-store/tests/segment_test.cpp @@ -980,58 +980,6 @@ TEST_F(SegmentTest, SharedNameIndexesSurviveReverseUnmountOrder) { } } -TEST_F(SegmentTest, SharedNameRegistrationsSurviveSegmentSnapshotRestore) { - SegmentManager source(BufferAllocatorType::OFFSET); - Segment first; - first.id = generate_uuid(); - first.name = "snapshot-shared-hostname"; - first.base = DEFAULT_CXL_BASE; - first.size = 64 * 1024 * 1024; - first.te_endpoint = "snapshot-first-endpoint"; - Segment second = first; - second.id = generate_uuid(); - second.base += first.size; - second.te_endpoint = "snapshot-second-endpoint"; - const UUID first_owner = generate_uuid(); - const UUID second_owner = generate_uuid(); - - { - auto segment_access = source.getSegmentAccess(); - ASSERT_EQ( - segment_access.MountSegment(first, first_owner, client_liveness_), - ErrorCode::OK); - ASSERT_EQ(segment_access.MountSegment( - second, second_owner, - std::make_shared( - ClientLivenessRecord::Clock::now())), - ErrorCode::OK); - } - auto serialized = - SegmentSerializer(&source).Serialize(LocalSsdPersistedState{}); - ASSERT_TRUE(serialized.has_value()); - - SegmentManager restored(BufferAllocatorType::OFFSET); - auto restored_state = SegmentSerializer(&restored).Deserialize(*serialized); - ASSERT_TRUE(restored_state.has_value()); - { - auto allocator_access = restored.getAllocatorAccess(); - const auto* registrations = - allocator_access.getAllocatorManager().getAllocators(first.name); - ASSERT_NE(registrations, nullptr); - EXPECT_EQ(registrations->size(), 2u); - } - - std::vector> segments; - ASSERT_EQ(restored.getSegmentAccess().GetAllSegments(segments), - ErrorCode::OK); - std::unordered_map> owners; - for (const auto& [segment, owner] : segments) { - owners.emplace(segment.id, owner); - } - EXPECT_EQ(owners.at(first.id), first_owner); - EXPECT_EQ(owners.at(second.id), second_owner); -} - TEST_F(SegmentTest, HostOrderedSegmentsRotateWithinSameHostByKey) { SegmentManager segment_manager; diff --git a/mooncake-store/tests/serializer_test.cpp b/mooncake-store/tests/serializer_test.cpp deleted file mode 100644 index 538d2ee05c..0000000000 --- a/mooncake-store/tests/serializer_test.cpp +++ /dev/null @@ -1,225 +0,0 @@ -#include -#include - -#include "serializer.h" -#include "serialize/serializer.h" -#include "segment.h" - -namespace mooncake::test { - -// Example class implementing serialization following the usage documentation -class ExampleClass { - public: - ExampleClass() : id_(0), value_(0.0), name_("") {} - - ExampleClass(int id, double value, const std::string& name) - : id_(id), value_(value), name_(name) {} - - // Serialization method (works with both counter and writer) - template - void serialize_to(T& serializer) const { - serializer.write(&id_, sizeof(id_)); - serializer.write(&value_, sizeof(value_)); - - // Serialize string length first, then string data - size_t name_length = name_.length(); - serializer.write(&name_length, sizeof(name_length)); - if (!name_.empty()) { - serializer.write(name_.data(), name_.length()); - } - } - - // Deserialization method - template - static std::shared_ptr deserialize_from(T& serializer) { - try { - auto obj = std::make_shared(); - - // Deserialize basic members - serializer.read(&obj->id_, sizeof(obj->id_)); - serializer.read(&obj->value_, sizeof(obj->value_)); - - // Deserialize string - size_t name_length; - serializer.read(&name_length, sizeof(name_length)); - if (name_length > 0) { - obj->name_.resize(name_length); - serializer.read(&obj->name_[0], name_length); - } - - return obj; - } catch (const std::exception& e) { - return nullptr; - } - } - - // Getters for testing - int getId() const { return id_; } - double getValue() const { return value_; } - const std::string& getName() const { return name_; } - - // Equality operator for testing - bool operator==(const ExampleClass& other) const { - return id_ == other.id_ && value_ == other.value_ && - name_ == other.name_; - } - - protected: - int id_; - double value_; - std::string name_; -}; - -class ExampleClassWithException { - public: - ExampleClassWithException() : value_(0) {} - ExampleClassWithException(int value) : value_(value) {} - - template - void serialize_to(T& serializer) const { - serializer.write(&value_, sizeof(value_)); - } - - template - static std::shared_ptr deserialize_from( - T& serializer) { - throw std::runtime_error("throw_exception"); - } - - private: - int value_; -}; - -class SerializerTest : public ::testing::Test { - protected: - void SetUp() override { - google::InitGoogleLogging("SerializerTest"); - FLAGS_logtostderr = true; - } - - void TearDown() override { google::ShutdownGoogleLogging(); } -}; - -TEST_F(SerializerTest, ExampleClassSerialization) { - // Create an example object - ExampleClass original(42, 3.14159, "Test Object"); - - // Test serialization - std::vector buffer; - ASSERT_EQ(serialize_to(original, buffer), ErrorCode::OK); - ASSERT_FALSE(buffer.empty()); - - // Test deserialization - auto restored = deserialize_from(buffer); - ASSERT_NE(restored, nullptr); - - // Verify the deserialized object matches the original - EXPECT_EQ(restored->getId(), original.getId()); - EXPECT_DOUBLE_EQ(restored->getValue(), original.getValue()); - EXPECT_EQ(restored->getName(), original.getName()); - EXPECT_TRUE(*restored == original); -} - -TEST_F(SerializerTest, ExampleClassSerializationWithSharedPtr) { - // Test with shared_ptr - auto original = - std::make_shared(777, 2.718, "Shared Pointer Test"); - - std::vector buffer; - ASSERT_EQ(serialize_to(original, buffer), ErrorCode::OK); - - auto restored = deserialize_from(buffer); - ASSERT_NE(restored, nullptr); - EXPECT_TRUE(*restored == *original); -} - -TEST_F(SerializerTest, ExampleClassSerializationNullPointer) { - // Test with null shared_ptr - std::shared_ptr null_ptr = nullptr; - - std::vector buffer; - ASSERT_EQ(serialize_to(null_ptr, buffer), ErrorCode::INVALID_PARAMS); -} - -TEST_F(SerializerTest, ExampleClassDeserializationCorruptedBuffer) { - // Create a valid object and serialize it - ExampleClass original(1, 1.0, "Test"); - std::vector buffer; - ASSERT_EQ(serialize_to(original, buffer), ErrorCode::OK); - - // Corrupt the buffer by removing the last byte - buffer.pop_back(); - - // Try to deserialize corrupted buffer - auto restored = deserialize_from(buffer); - EXPECT_EQ(restored, nullptr); -} - -TEST_F(SerializerTest, ExampleClassDeserializationWithException) { - // Create a valid object and serialize it - ExampleClassWithException original(1); - std::vector buffer; - ASSERT_EQ(serialize_to(original, buffer), ErrorCode::OK); - - // Try to deserialize the buffer, the deserialization method will throw an - // exception. - auto restored = deserialize_from(buffer); - EXPECT_EQ(restored, nullptr); -} - -TEST_F(SerializerTest, MountedSegmentSerializationPreservesHostId) { - MountedSegment original; - original.segment.id = generate_uuid(); - original.segment.name = "segment_host1"; - original.segment.base = 0x300000000; - original.segment.size = 1024 * 1024; - original.segment.te_endpoint = "segment_host1"; - original.segment.host_id = "host1"; - original.status = SegmentStatus::OK; - - msgpack::sbuffer buffer; - MsgpackPacker packer(&buffer); - ASSERT_TRUE( - Serializer::serialize(original, packer).has_value()); - - auto object_handle = msgpack::unpack(buffer.data(), buffer.size()); - auto restored = - Serializer::deserialize(object_handle.get()); - ASSERT_TRUE(restored.has_value()); - EXPECT_EQ(restored->segment.id, original.segment.id); - EXPECT_EQ(restored->segment.name, original.segment.name); - EXPECT_EQ(restored->segment.host_id, original.segment.host_id); - EXPECT_EQ(restored->status, original.status); -} - -TEST_F(SerializerTest, MountedSegmentDeserializesLegacyFormatWithoutHostId) { - const UUID segment_id = generate_uuid(); - - msgpack::sbuffer buffer; - MsgpackPacker packer(&buffer); - packer.pack_array(8); - packer.pack(UuidToString(segment_id)); - packer.pack(std::string("legacy_segment")); - packer.pack(static_cast(0x300000000)); - packer.pack(static_cast(1024 * 1024)); - packer.pack(std::string("legacy_segment")); - packer.pack(static_cast(SegmentStatus::OK)); - packer.pack(false); - packer.pack_nil(); - - auto object_handle = msgpack::unpack(buffer.data(), buffer.size()); - auto restored = - Serializer::deserialize(object_handle.get()); - ASSERT_TRUE(restored.has_value()); - EXPECT_EQ(restored->segment.id, segment_id); - EXPECT_EQ(restored->segment.name, "legacy_segment"); - EXPECT_TRUE(restored->segment.host_id.empty()); - EXPECT_EQ(restored->status, SegmentStatus::OK); -} - -} // namespace mooncake::test - -int main(int argc, char** argv) { - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/mooncake-store/tests/task_manager_test.cpp b/mooncake-store/tests/task_manager_test.cpp index 7c935f4f06..7ce657a07f 100644 --- a/mooncake-store/tests/task_manager_test.cpp +++ b/mooncake-store/tests/task_manager_test.cpp @@ -269,154 +269,6 @@ TEST_F(ClientTaskManagerTest, PruneExpiredTasksProcessingTimeoutFreesSlot) { EXPECT_EQ(second[0].status, TaskStatus::PROCESSING); } -TEST_F(ClientTaskManagerTest, SerializerRoundTrip) { - ClientTaskManager manager({10000, 10000, 10000, 0, 0, 0}); - UUID client_id1 = generate_uuid(); - UUID client_id2 = generate_uuid(); - - // Create task - auto t1 = - manager.get_write_access().submit_task_typed( - client_id1, - ReplicaCopyPayload{ - .key = "pending_key", .source = "", .targets = {"seg1"}}); - auto t2 = - manager.get_write_access().submit_task_typed( - client_id2, - ReplicaCopyPayload{ - .key = "processing_key", .source = "", .targets = {"seg2"}}); - auto t3 = - manager.get_write_access().submit_task_typed( - client_id1, ReplicaMovePayload{ - .key = "key3", .source = "seg1", .target = "seg4"}); - auto t4 = - manager.get_write_access().submit_task_typed( - client_id2, ReplicaMovePayload{ - .key = "key4", .source = "seg2", .target = "seg5"}); - - ASSERT_TRUE(t1.has_value()); - ASSERT_TRUE(t2.has_value()); - ASSERT_TRUE(t3.has_value()); - ASSERT_TRUE(t4.has_value()); - - // Pop t1 from client_id1 to make task processing - manager.get_write_access().pop_tasks(client_id1, 1); - - // Pop t2 t4 from client_id2 to make task processing - manager.get_write_access().pop_tasks(client_id2, 2); - - // Complete t1 as SUCCESS, t2 as FAILED - std::ignore = manager.get_write_access().complete_task( - client_id1, t1.value(), TaskStatus::SUCCESS, "Done"); - std::ignore = manager.get_write_access().complete_task( - client_id2, t2.value(), TaskStatus::FAILED, "Failed"); - - // Get original tasks - auto task1_before = manager.get_read_access().find_task_by_id(t1.value()); - auto task2_before = manager.get_read_access().find_task_by_id(t2.value()); - auto task3_before = manager.get_read_access().find_task_by_id(t3.value()); - auto task4_before = manager.get_read_access().find_task_by_id(t4.value()); - - // Serialize - TaskManagerSerializer serializer(&manager); - auto serialized = serializer.Serialize(); - ASSERT_TRUE(serialized.has_value()); - - // Deserialize into new manager - ClientTaskManager manager2({10000, 10000, 10000, 0, 0, 0}); - TaskManagerSerializer serializer2(&manager2); - auto result = serializer2.Deserialize(serialized.value()); - ASSERT_TRUE(result.has_value()); - - // Verify statuses - auto task1 = manager2.get_read_access().find_task_by_id(t1.value()); - auto task2 = manager2.get_read_access().find_task_by_id(t2.value()); - auto task3 = manager2.get_read_access().find_task_by_id(t3.value()); - auto task4 = manager2.get_read_access().find_task_by_id(t4.value()); - - ASSERT_TRUE(task1.has_value()); - ASSERT_TRUE(task2.has_value()); - ASSERT_TRUE(task3.has_value()); - ASSERT_TRUE(task4.has_value()); - - auto time_diff = [](auto time1, auto time2) { - return std::abs( - std::chrono::duration_cast(time1 - time2) - .count()); - }; - - EXPECT_EQ(task1->id, task1_before->id); - EXPECT_EQ(task1->type, TaskType::REPLICA_COPY); - EXPECT_EQ(task1->status, TaskStatus::SUCCESS); - EXPECT_FALSE(task1->payload.empty()); - EXPECT_EQ(task1->assigned_client, client_id1); - EXPECT_LE(time_diff(task1->created_at, task1_before->created_at), 1); - EXPECT_LE(time_diff(task1->last_updated_at, task1_before->last_updated_at), - 1); - EXPECT_EQ(task2->id, task2_before->id); - EXPECT_EQ(task2->type, TaskType::REPLICA_COPY); - EXPECT_EQ(task2->status, TaskStatus::FAILED); - EXPECT_LE(time_diff(task2->created_at, task2_before->created_at), 1); - EXPECT_LE(time_diff(task2->last_updated_at, task2_before->last_updated_at), - 1); - EXPECT_FALSE(task2->payload.empty()); - EXPECT_EQ(task2->assigned_client, client_id2); - EXPECT_EQ(task3->id, task3_before->id); - EXPECT_EQ(task3->type, TaskType::REPLICA_MOVE); - EXPECT_EQ(task3->status, TaskStatus::PENDING); - EXPECT_LE(time_diff(task3->created_at, task3_before->created_at), 1); - EXPECT_LE(time_diff(task3->last_updated_at, task3_before->last_updated_at), - 1); - EXPECT_FALSE(task3->payload.empty()); - EXPECT_EQ(task3->assigned_client, client_id1); - EXPECT_EQ(task4->id, task4_before->id); - EXPECT_EQ(task4->type, TaskType::REPLICA_MOVE); - EXPECT_EQ(task4->status, TaskStatus::PROCESSING); - EXPECT_LE(time_diff(task4->created_at, task4_before->created_at), 1); - EXPECT_LE(time_diff(task4->last_updated_at, task4_before->last_updated_at), - 1); - EXPECT_FALSE(task4->payload.empty()); - EXPECT_EQ(task4->assigned_client, client_id2); -} - -TEST_F(ClientTaskManagerTest, SerializerEmptyManager) { - ClientTaskManager manager({10000, 10000, 10000, 0, 0, 0}); - - TaskManagerSerializer serializer(&manager); - auto serialized = serializer.Serialize(); - ASSERT_TRUE(serialized.has_value()); - - ClientTaskManager manager2({10000, 10000, 10000, 0, 0, 0}); - TaskManagerSerializer serializer2(&manager2); - auto result = serializer2.Deserialize(serialized.value()); - ASSERT_TRUE(result.has_value()); - - // Verify empty - auto tasks = manager2.get_read_access(); - EXPECT_EQ(tasks.size(), 0u); -} - -TEST_F(ClientTaskManagerTest, SerializerReset) { - ClientTaskManager manager({10000, 10000, 10000, 0, 0, 0}); - UUID client_id = generate_uuid(); - - auto t1 = - manager.get_write_access().submit_task_typed( - client_id, ReplicaCopyPayload{ - .key = "key1", .source = "", .targets = {"seg1"}}); - ASSERT_TRUE(t1.has_value()); - - TaskManagerSerializer serializer(&manager); - serializer.Reset(); - - // Verify all data cleared - auto task = manager.get_read_access().find_task_by_id(t1.value()); - EXPECT_FALSE(task.has_value()); - - auto tasks = manager.get_read_access(); - EXPECT_EQ(tasks.size(), 0u); -} - TEST_F(ClientTaskManagerTest, FailPendingTaskSkipsQueuedAssignment) { ClientTaskManager manager({10000, 10000, 10000, 0, 0, 3}); UUID client_id = generate_uuid();