diff --git a/be/src/storage/index/primary_key_index.cpp b/be/src/storage/index/primary_key_index.cpp index 28fa377e4b4aee..080e19d503b727 100644 --- a/be/src/storage/index/primary_key_index.cpp +++ b/be/src/storage/index/primary_key_index.cpp @@ -112,12 +112,16 @@ Status PrimaryKeyIndexReader::parse_index(io::FileReaderSPtr file_reader, OlapReaderStatistics* pk_index_load_stats, const io::IOContext* source_io_ctx) { // parse primary key index - _index_reader.reset(new segment_v2::IndexedColumnReader(file_reader, meta.primary_key_index())); - _index_reader->set_is_pk_index(true); + auto index_reader = std::make_unique(file_reader, + meta.primary_key_index()); + index_reader->set_is_pk_index(true); auto io_ctx = create_index_io_context(source_io_ctx, pk_index_load_stats); - RETURN_IF_ERROR(_index_reader->load(!config::disable_pk_storage_page_cache, false, - pk_index_load_stats, &io_ctx)); + RETURN_IF_ERROR(index_reader->load(!config::disable_pk_storage_page_cache, false, + pk_index_load_stats, &io_ctx)); + // A failed initialization must not leave partially loaded root pages pinned in + // a cached segment without a corresponding eviction charge. + _index_reader = std::move(index_reader); _index_parsed = true; return Status::OK(); } diff --git a/be/src/storage/segment/segment.cpp b/be/src/storage/segment/segment.cpp index 8978d7e01eac03..3b7bd70b851d35 100644 --- a/be/src/storage/segment/segment.cpp +++ b/be/src/storage/segment/segment.cpp @@ -783,7 +783,11 @@ Status Segment::load_index(OlapReaderStatistics* stats, const io::IOContext* sou _pk_index_reader = std::make_unique(); RETURN_IF_ERROR(_pk_index_reader->parse_index(_file_reader, *_pk_index_meta, stats, source_io_ctx)); - // _meta_mem_usage += _pk_index_reader->get_memory_size(); + _pk_index_cache_bytes.store(_pk_index_reader->get_memory_size(), + std::memory_order_relaxed); + if (_cache_charge_callback) { + _cache_charge_callback(cache_charge()); + } return Status::OK(); } else { // read and parse short key index page diff --git a/be/src/storage/segment/segment.h b/be/src/storage/segment/segment.h index a39114187c2021..31715425599ac0 100644 --- a/be/src/storage/segment/segment.h +++ b/be/src/storage/segment/segment.h @@ -22,7 +22,9 @@ #include #include +#include #include +#include #include #include // for unique_ptr #include @@ -181,6 +183,18 @@ class Segment : public std::enable_shared_from_this, public MetadataAdd // another method `get_metadata_size` not include the column reader, only the segment object itself. int64_t meta_mem_usage() const { return _meta_mem_usage; } + // PK pages are tracked by PKIndexPageCache, but their pinned bytes must also + // count toward SegmentCache eviction. Do not add these bytes to metadata tracking. + size_t cache_charge() const { + return _meta_mem_usage + _pk_index_cache_bytes.load(std::memory_order_relaxed); + } + + // Install before publishing this segment in SegmentCache; immutable afterwards. + void set_cache_charge_callback(std::function callback) { + DCHECK(!_cache_charge_callback); + _cache_charge_callback = std::move(callback); + } + // Variant paths use segment metadata; other columns use `read_type`. std::shared_ptr get_data_type_of(const TabletColumn& read_column, const DataTypePtr& read_type, @@ -292,6 +306,8 @@ class Segment : public std::enable_shared_from_this, public MetadataAdd // The memory consumed by querying is tracked in segment iterator. int64_t _meta_mem_usage; int64_t _tracked_meta_mem_usage = 0; + std::atomic _pk_index_cache_bytes {0}; + std::function _cache_charge_callback; RowsetId _rowset_id; TabletSchemaSPtr _tablet_schema; diff --git a/be/src/storage/segment/segment_loader.cpp b/be/src/storage/segment/segment_loader.cpp index e1c12cbf486131..c7f4a1a6d3b2ef 100644 --- a/be/src/storage/segment/segment_loader.cpp +++ b/be/src/storage/segment/segment_loader.cpp @@ -43,8 +43,27 @@ bool SegmentCache::lookup(const SegmentCache::CacheKey& key, SegmentCacheHandle* void SegmentCache::insert(const SegmentCache::CacheKey& key, SegmentCache::CacheValue& value, SegmentCacheHandle* handle) { + // Keep the cache weakly referenced: the cache owns the segment. The raw segment + // pointer is only compared while its own initialization invokes the callback. + value.segment->set_cache_charge_callback([cache = std::weak_ptr(_cache), + encoded_key = key.encode(), + segment = value.segment.get()](size_t charge) { + auto owner = cache.lock(); + if (!owner) { + return; // The Segment may outlive the cache during shutdown. + } + auto* entry = owner->lookup(encoded_key); + if (entry == nullptr) { + return; // Already evicted. + } + // Another instance of this segment may have replaced the cached entry. + if (static_cast(owner->value(entry))->segment.get() == segment) { + owner->update_charge(entry, charge); + } + owner->release(entry); + }); auto* lru_handle = - LRUCachePolicy::insert(key.encode(), &value, value.segment->meta_mem_usage(), + LRUCachePolicy::insert(key.encode(), &value, value.segment->cache_charge(), value.segment->meta_mem_usage(), CachePriority::NORMAL); handle->push_segment(this, lru_handle); } diff --git a/be/src/util/lru_cache.cpp b/be/src/util/lru_cache.cpp index 353c7303212376..5d44236c9a592e 100644 --- a/be/src/util/lru_cache.cpp +++ b/be/src/util/lru_cache.cpp @@ -335,6 +335,48 @@ Cache::Handle* LRUCache::lookup(const CacheKey& key, uint32_t hash) { return reinterpret_cast(e); } +void LRUCache::update_charge(Cache::Handle* handle, size_t charge) { + auto* e = reinterpret_cast(handle); + LRUHandle* to_remove_head = nullptr; + { + std::lock_guard l(_mutex); + // Erase/replacement may have removed this entry while its handle was held. + if (!e->in_cache || charge <= e->charge) { + return; + } + DCHECK_GT(e->refs, 1); + size_t delta = charge - e->charge; + e->charge = charge; + e->total_size += delta; + _usage += delta; + if (e->total_size > _capacity) { + // This entry cannot fit even after evicting every other entry. Drop + // cache ownership now, while outstanding handles keep it alive. + bool removed = _table.remove(e); + DCHECK(removed); + e->in_cache = false; + _unref(e); + _usage -= e->total_size; + } + // Growing an existing entry does not need another element slot. The + // insertion eviction helpers also check count >= limit, so only call + // them when the updated usage exceeds the byte capacity. + if (_usage > _capacity) { + if (_cache_value_check_timestamp) { + _evict_from_lru_with_time(0, &to_remove_head); + } else { + _evict_from_lru(0, &to_remove_head); + } + } + } + // Cache values can release other caches; never destroy them under the shard lock. + while (to_remove_head != nullptr) { + LRUHandle* next = to_remove_head->next; + to_remove_head->free(); + to_remove_head = next; + } +} + void LRUCache::release(Cache::Handle* handle) { if (handle == nullptr) { return; @@ -769,6 +811,11 @@ void ShardedLRUCache::release(Handle* handle) { _shards[_shard(h->hash)]->release(handle); } +void ShardedLRUCache::update_charge(Handle* handle, size_t charge) { + auto* h = reinterpret_cast(handle); + _shards[_shard(h->hash)]->update_charge(handle, charge); +} + void ShardedLRUCache::erase(const CacheKey& key) { const uint32_t hash = _hash_slice(key); _shards[_shard(hash)]->erase(key, hash); diff --git a/be/src/util/lru_cache.h b/be/src/util/lru_cache.h index f01fd0cd6cf5e2..718baabba7f275 100644 --- a/be/src/util/lru_cache.h +++ b/be/src/util/lru_cache.h @@ -193,6 +193,10 @@ class Cache { // REQUIRES: handle must have been returned by a method on *this. virtual void release(Handle* handle) = 0; + // Grow an entry's eviction charge after lazy initialization, without changing its + // memory tracker. The caller must hold a live handle. Stale smaller charges are ignored. + virtual void update_charge(Handle* handle, size_t charge) = 0; + // Return the value encapsulated in a handle returned by a // successful lookup(). // REQUIRES: handle must not have been released yet. @@ -348,6 +352,7 @@ class LRUCache { CachePriority priority = CachePriority::NORMAL); Cache::Handle* lookup(const CacheKey& key, uint32_t hash); void release(Cache::Handle* handle); + void update_charge(Cache::Handle* handle, size_t charge); void erase(const CacheKey& key, uint32_t hash); PrunedInfo prune(); PrunedInfo prune_if(CachePrunePredicate pred, bool lazy_mode = false); @@ -420,6 +425,7 @@ class ShardedLRUCache : public Cache { CachePriority priority = CachePriority::NORMAL) override; Handle* lookup(const CacheKey& key) override; void release(Handle* handle) override; + void update_charge(Handle* handle, size_t charge) override; void erase(const CacheKey& key) override; void* value(Handle* handle) override; uint64_t new_id() override; @@ -484,6 +490,7 @@ class DummyLRUCache : public Cache { CachePriority priority = CachePriority::NORMAL) override; Handle* lookup(const CacheKey& key) override { return nullptr; }; void release(Handle* handle) override; + void update_charge(Handle* handle, size_t charge) override {} void erase(const CacheKey& key) override {}; void* value(Handle* handle) override; uint64_t new_id() override { return 0; }; diff --git a/be/test/storage/cache/lru_cache_test.cpp b/be/test/storage/cache/lru_cache_test.cpp index a31d532a5a4686..5c248838ad65e7 100644 --- a/be/test/storage/cache/lru_cache_test.cpp +++ b/be/test/storage/cache/lru_cache_test.cpp @@ -21,6 +21,7 @@ #include #include +#include #include #include @@ -297,6 +298,180 @@ TEST_F(CacheTest, EvictionPolicyWithDurable) { EXPECT_EQ(201, Lookup(200)); } +TEST_F(CacheTest, UpdateChargeDoesNotDoubleTrackMemory) { + init_size_cache(1024 * 1024); + const CacheKey key("growing"); + auto* handle = cache()->insert(key, new CacheValue(EncodeValue(1)), 100, 100); + const auto usage = cache()->get_usage(); + const auto tracked = cache()->mem_consumption(); + cache()->_cache->update_charge(handle, 500); + EXPECT_EQ(cache()->get_usage(), usage + 400); + EXPECT_EQ(cache()->mem_consumption(), tracked); + // Repeated initialization and older snapshots cannot undo the larger charge. + cache()->_cache->update_charge(handle, 500); + cache()->_cache->update_charge(handle, 200); + EXPECT_EQ(cache()->get_usage(), usage + 400); + cache()->_cache->update_charge(handle, 2 * 1024 * 1024); + EXPECT_EQ(cache()->get_usage(), 0); + // Detaching an oversized entry does not release its live handle's memory. + EXPECT_EQ(cache()->mem_consumption(), tracked); + cache()->release(handle); + cache()->prune_all(true); + EXPECT_EQ(cache()->get_usage(), 0); + EXPECT_EQ(cache()->mem_consumption(), 0); +} + +TEST_F(CacheTest, UpdateChargeEvictsIdleEntriesWhenEntryStillFits) { + LRUCache cache(LRUCacheType::SIZE); + cache.set_capacity(1024); + CacheKey idle("idle"), durable("durable"), growing("growing"); + cache.release(cache.insert(idle, 1, new CacheValue(EncodeValue(1)), 100)); + cache.release( + cache.insert(durable, 3, new CacheValue(EncodeValue(3)), 100, CachePriority::DURABLE)); + auto* handle = cache.insert(growing, 2, new CacheValue(EncodeValue(2)), 100); + cache.update_charge(handle, 700); + EXPECT_EQ(cache.lookup(idle, 1), nullptr); + EXPECT_LE(cache.get_usage(), 1024); + auto* retained = cache.lookup(durable, 3); + ASSERT_NE(retained, nullptr); + cache.release(retained); + cache.release(handle); + EXPECT_EQ(cache.get_element_count(), 2); +} + +TEST_F(CacheTest, UpdateChargeAtElementLimitOnlyEvictsAboveByteCapacity) { + for (bool check_timestamp : {false, true}) { + LRUCache cache(LRUCacheType::SIZE, true); + cache.set_capacity(1024); + cache.set_element_count_capacity(2); + cache.set_cache_value_time_extractor([](const void* value) -> int64_t { + return DecodeValue(static_cast(value)->value); + }); + cache.set_cache_value_check_timestamp(check_timestamp); + CacheKey idle("idle"), growing("growing"); + cache.release(cache.insert(idle, 1, new CacheValue(EncodeValue(1)), 100)); + auto* handle = cache.insert(growing, 2, new CacheValue(EncodeValue(2)), 100); + ASSERT_EQ(cache.get_element_count(), 2); + const auto usage = cache.get_usage(); + + cache.update_charge(handle, 200); + EXPECT_EQ(cache.get_usage(), usage + 100); + EXPECT_EQ(cache.get_element_count(), 2); + auto* retained = cache.lookup(idle, 1); + ASSERT_NE(retained, nullptr); + cache.release(retained); + + // Exactly filling byte capacity must also preserve both entries. + const auto full_charge = 100 + cache.get_capacity() - usage; + cache.update_charge(handle, full_charge); + EXPECT_EQ(cache.get_usage(), cache.get_capacity()); + EXPECT_EQ(cache.get_element_count(), 2); + + // Crossing byte capacity still evicts the idle entry normally. + cache.update_charge(handle, full_charge + 1); + EXPECT_EQ(cache.lookup(idle, 1), nullptr); + EXPECT_EQ(cache.get_element_count(), 1); + EXPECT_LE(cache.get_usage(), cache.get_capacity()); + cache.release(handle); + auto* growing_handle = cache.lookup(growing, 2); + ASSERT_NE(growing_handle, nullptr); + cache.release(growing_handle); + } +} + +TEST_F(CacheTest, UpdateChargeOversizedEntryPreservesOtherEntries) { + for (bool check_timestamp : {false, true}) { + LRUCache cache(LRUCacheType::SIZE, true); + cache.set_capacity(1024); + cache.set_cache_value_time_extractor([](const void* value) -> int64_t { + return DecodeValue(static_cast(value)->value); + }); + cache.set_cache_value_check_timestamp(check_timestamp); + CacheKey idle("idle"), durable("durable"), growing("growing"); + cache.release(cache.insert(idle, 1, new CacheValue(EncodeValue(1)), 100)); + cache.release(cache.insert(durable, 3, new CacheValue(EncodeValue(3)), 100, + CachePriority::DURABLE)); + const auto usage = cache.get_usage(); + auto* handle = cache.insert(growing, 2, new CacheValue(EncodeValue(2)), 100); + cache.update_charge(handle, 2048); + EXPECT_EQ(cache.get_usage(), usage); + EXPECT_EQ(cache.get_element_count(), 2); + EXPECT_EQ(cache.lookup(growing, 2), nullptr); + // The detached entry remains usable through the caller's live handle. + EXPECT_EQ(DecodeValue(static_cast(reinterpret_cast(handle)->value) + ->value), + 2); + cache.release(handle); + auto* idle_handle = cache.lookup(idle, 1); + ASSERT_NE(idle_handle, nullptr); + cache.release(idle_handle); + auto* durable_handle = cache.lookup(durable, 3); + ASSERT_NE(durable_handle, nullptr); + cache.release(durable_handle); + EXPECT_EQ(cache.get_usage(), usage); + } +} + +TEST_F(CacheTest, UpdateChargeIgnoresReplacedEntry) { + LRUCache cache(LRUCacheType::SIZE); + cache.set_capacity(4096); + CacheKey key("replaced"); + auto* old = cache.insert(key, 1, new CacheValue(EncodeValue(1)), 100); + auto* current = cache.insert(key, 1, new CacheValue(EncodeValue(2)), 200); + const auto usage = cache.get_usage(); + cache.update_charge(old, 8192); + EXPECT_EQ(cache.get_usage(), usage); + cache.release(old); + cache.release(current); + EXPECT_EQ(cache.get_usage(), usage); +} + +TEST_F(CacheTest, UpdateChargeConcurrentGrowth) { + LRUCache cache(LRUCacheType::SIZE); + cache.set_capacity(1024 * 1024); + CacheKey key("concurrent"); + auto* handle = cache.insert(key, 1, new CacheValue(EncodeValue(1)), 100); + const auto usage = cache.get_usage(); + std::vector> updates; + for (size_t charge = 200; charge <= 900; charge += 100) { + auto* held = cache.lookup(key, 1); + ASSERT_NE(held, nullptr); + updates.emplace_back(std::async(std::launch::async, [&, held, charge] { + cache.update_charge(held, charge); + cache.release(held); + })); + } + for (auto& update : updates) { + update.get(); + } + EXPECT_EQ(cache.get_usage(), usage + 800); + cache.release(handle); +} + +TEST_F(CacheTest, UpdateChargePreservesOtherOutstandingHandles) { + LRUCache cache(LRUCacheType::SIZE); + cache.set_capacity(1024); + CacheKey key("shared"); + auto* first = cache.insert(key, 1, new CacheValueWithKey(1, EncodeValue(1)), 100); + auto* second = cache.lookup(key, 1); + ASSERT_NE(second, nullptr); + cache.update_charge(first, 2048); + EXPECT_EQ(cache.get_usage(), 0); + EXPECT_EQ(cache.lookup(key, 1), nullptr); + cache.release(first); + EXPECT_TRUE(_deleted_keys.empty()); + EXPECT_EQ( + DecodeValue(static_cast(reinterpret_cast(second)->value) + ->value), + 1); + cache.update_charge(second, 4096); + EXPECT_EQ(cache.get_usage(), 0); + cache.release(second); + ASSERT_EQ(_deleted_keys.size(), 1); + EXPECT_EQ(_deleted_keys.front(), 1); + EXPECT_EQ(cache.get_usage(), 0); +} + TEST_F(CacheTest, Usage) { LRUCache cache(LRUCacheType::SIZE); cache.set_capacity(1040); diff --git a/be/test/storage/segment/segment_primary_key_cache_test.cpp b/be/test/storage/segment/segment_primary_key_cache_test.cpp new file mode 100644 index 00000000000000..7de4aa4625b27f --- /dev/null +++ b/be/test/storage/segment/segment_primary_key_cache_test.cpp @@ -0,0 +1,190 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include + +#include +#include +#include +#include + +#include "core/field.h" +#include "io/fs/local_file_system.h" +#include "runtime/exec_env.h" +#include "storage/index/primary_key_index.h" +#include "storage/segment/segment.h" +#include "storage/segment/segment_loader.h" +#include "storage/segment/vertical_segment_writer.h" +#include "storage/storage_engine.h" +#include "storage/tablet/tablet_schema_helper.h" + +namespace doris { + +TabletSchemaSPtr create_schema(const std::vector& columns, KeysType keys_type); +using Generator = std::function; +void build_segment(VerticalSegmentWriterOptions opts, TabletSchemaSPtr build_schema, + size_t segment_id, TabletSchemaSPtr query_schema, size_t nrows, + Generator generator, std::shared_ptr* res, std::string segment_dir); + +class SegmentPrimaryKeyCacheTest : public testing::Test { +protected: + void SetUp() override { + auto fs = io::global_local_filesystem(); + ASSERT_TRUE(fs->delete_directory(_dir).ok()); + ASSERT_TRUE(fs->create_directory(_dir).ok()); + ExecEnv::GetInstance()->set_storage_engine( + std::make_unique(EngineOptions {})); + _schema = create_schema( + {create_varchar_key(0, false), + create_int_value(1, FieldAggregationMethod::OLAP_FIELD_AGGREGATION_REPLACE, + false)}, + UNIQUE_KEYS); + VerticalSegmentWriterOptions opts; + opts.enable_unique_key_merge_on_write = true; + auto generator = [](size_t rid, int cid, Field& field) { + if (cid == 0) { + // Each key exceeds the normal PK data-page target, producing a large value index. + field = Field::create_field(std::string(50000, 'a' + rid)); + } else { + field = Field::create_field(static_cast(rid)); + } + }; + build_segment(opts, _schema, 0, _schema, 8, generator, &_segment, _dir); + ASSERT_NE(_segment, nullptr); + ASSERT_EQ(_segment->_pk_index_reader, nullptr); + ASSERT_FALSE(_segment->_load_index_once.has_called()); + ASSERT_FALSE(_segment->_load_pk_bf_once.has_called()); + ASSERT_FALSE(_segment->_pk_index_meta->primary_key_index() + .value_index_meta() + .is_root_data_page()); + _present_key = _segment->min_key(); + + // Find a definite BF miss without initializing the segment's own reader. Do not assume + // that an arbitrary absent key is rejected, because bloom filters allow false positives. + PrimaryKeyIndexReader probe; + ASSERT_TRUE( + probe.parse_bf(_segment->file_reader(), *_segment->_pk_index_meta, nullptr).ok()); + for (int i = 0; i < 10000; ++i) { + auto candidate = _present_key + std::to_string(i); + if (!probe.check_present(Slice(candidate))) { + _missing_key = std::move(candidate); + break; + } + } + ASSERT_FALSE(_missing_key.empty()); + } + + void TearDown() override { + _segment.reset(); + ExecEnv::GetInstance()->set_storage_engine(nullptr); + EXPECT_TRUE(io::global_local_filesystem()->delete_directory(_dir).ok()); + } + + Status lookup(const std::string& key, RowLocation* location) { + return _segment->lookup_row_key(Slice(key), _schema.get(), false, false, location, nullptr); + } + + const std::string _dir = "./ut_dir/segment_primary_key_cache_test"; + TabletSchemaSPtr _schema; + std::shared_ptr _segment; + std::string _present_key; + std::string _missing_key; +}; + +TEST_F(SegmentPrimaryKeyCacheTest, CachedRootPagesIncreaseChargeOnlyOnce) { + SegmentCache cache(1024 * 1024 * 1024, 10000); + SegmentCacheHandle handle; + const SegmentCache::CacheKey key(_segment->rowset_id(), _segment->id()); + cache.insert(key, *new SegmentCache::CacheValue(_segment), &handle); + const auto usage = cache.get_usage(); + const auto tracked = cache.mem_consumption(); + RowLocation location; + ASSERT_TRUE(lookup(_missing_key, &location).is()); + // Keep the original index-first lookup: even a definite BF miss loads the roots. + const auto roots = _segment->_pk_index_reader->get_memory_size(); + EXPECT_GT(roots, 300000); + EXPECT_EQ(cache.get_usage(), usage + roots); + EXPECT_EQ(cache.mem_consumption(), tracked); + ASSERT_TRUE(lookup(_present_key, &location).ok()); + ASSERT_TRUE(_segment->load_index(nullptr).ok()); + EXPECT_EQ(cache.get_usage(), usage + roots); +} + +TEST_F(SegmentPrimaryKeyCacheTest, LazyRootLoadingEvictsOversizedSegment) { + // SegmentCache has 64 shards. Metadata fits; the large PK roots do not. + SegmentCache cache((_segment->cache_charge() + 65536) * 64, 10000); + SegmentCacheHandle handle; + const SegmentCache::CacheKey key(_segment->rowset_id(), _segment->id()); + cache.insert(key, *new SegmentCache::CacheValue(_segment), &handle); + ASSERT_EQ(cache.get_element_count(), 1); + ASSERT_TRUE(_segment->load_index(nullptr).ok()); + EXPECT_EQ(cache.get_element_count(), 0); + EXPECT_EQ(cache.get_usage(), 0); + // Eviction drops cache ownership, not the active operation's shared_ptr. + RowLocation location; + ASSERT_TRUE(lookup(_present_key, &location).ok()); + EXPECT_EQ(location.row_id, 0); +} + +TEST_F(SegmentPrimaryKeyCacheTest, EagerRootLoadingChargesOnInsert) { + const auto capacity = (_segment->cache_charge() + 65536) * 64; + ASSERT_TRUE(_segment->load_index(nullptr).ok()); + SegmentCache cache(capacity, 10000); + SegmentCacheHandle handle; + const SegmentCache::CacheKey key(_segment->rowset_id(), _segment->id()); + cache.insert(key, *new SegmentCache::CacheValue(_segment), &handle); + EXPECT_EQ(cache.get_usage(), 0); + EXPECT_EQ(cache.get_element_count(), 0); + EXPECT_EQ(handle.get_segments().front(), _segment); +} + +TEST_F(SegmentPrimaryKeyCacheTest, OldSegmentDoesNotChargeItsReplacement) { + SegmentCache cache(1024 * 1024 * 1024, 10000); + const SegmentCache::CacheKey key(_segment->rowset_id(), _segment->id()); + SegmentCacheHandle old_handle, new_handle; + cache.insert(key, *new SegmentCache::CacheValue(_segment), &old_handle); + auto replacement = std::make_shared(_segment->id(), _segment->rowset_id(), + _schema, InvertedIndexFileInfo {}); + cache.insert(key, *new SegmentCache::CacheValue(replacement), &new_handle); + const auto usage = cache.get_usage(); + ASSERT_TRUE(_segment->load_index(nullptr).ok()); + EXPECT_EQ(cache.get_usage(), usage); + EXPECT_EQ(cache.get_element_count(), 1); +} + +TEST_F(SegmentPrimaryKeyCacheTest, SegmentCanOutliveCache) { + { + SegmentCache cache(1024 * 1024 * 1024, 10000); + SegmentCacheHandle handle; + const SegmentCache::CacheKey key(_segment->rowset_id(), _segment->id()); + cache.insert(key, *new SegmentCache::CacheValue(_segment), &handle); + } + ASSERT_TRUE(_segment->load_index(nullptr).ok()); + RowLocation location; + ASSERT_TRUE(lookup(_present_key, &location).ok()); +} + +TEST_F(SegmentPrimaryKeyCacheTest, FailedIndexLoadDoesNotPinPartialRoots) { + auto meta = *_segment->_pk_index_meta; + meta.mutable_primary_key_index()->mutable_value_index_meta()->mutable_root_page()->set_offset( + uint64_t {1} << 50); + PrimaryKeyIndexReader reader; + EXPECT_FALSE(reader.parse_index(_segment->file_reader(), meta, nullptr).ok()); + EXPECT_EQ(reader._index_reader, nullptr); +} + +} // namespace doris