This is an automated email from the ASF dual-hosted git repository.

gavinchou pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git


The following commit(s) were added to refs/heads/master by this push:
     new 7271798202a [fix](mow) Account for pinned PK roots in segment cache 
eviction (#68304)
7271798202a is described below

commit 7271798202adbb94dabcaa4af73bcf2e9fc4a76e
Author: hui lai <[email protected]>
AuthorDate: Fri Oct 16 13:23:59 2026 +0800

    [fix](mow) Account for pinned PK roots in segment cache eviction (#68304)
    
    ### What problem does this PR solve?
    
    SegmentCache retained large PK index roots while its LRU charge only
    reflected
    estimated metadata. With long keys, a metadata-sized cache entry can pin
    hundreds
    of MiB of decompressed index pages and prevent the PK page cache from
    reclaiming
    them.
    
    The ownership chain is:
    
    ```text
    SegmentCache
      └─ Segment
           └─ PrimaryKeyIndexReader (_pk_index_reader)
                └─ IndexedColumnReader (_index_reader)
                     └─ PageHandle
                          └─ PK root page memory ← also cached by 
PKIndexPageCache
    ```
    
    The reader's `PageHandle` and PKIndexPageCache refer to the same page
    allocation,
    not separate copies. Keeping a Segment in SegmentCache can therefore
    keep its
    loaded PK roots pinned even after the query finishes. Evicting the
    SegmentCache
    entry removes cache ownership; once all remaining Segment references are
    released,
    the reader releases its page handles and the PK page cache can reclaim
    those pages.
    
    Charge loaded PK roots to SegmentCache eviction, both at insertion and
    when a
    cached Segment initializes its PK index later. Physical memory remains
    tracked
    by the PK page allocator; the additional eviction weight is not counted
    again
    by MemTracker. Updates use the shard lock and check Segment identity so
    an old
    instance cannot charge its replacement. The callback holds only a weak
    cache
    reference. Failed index initialization discards partial readers and
    their page
    handles.
    
    The original eager rowset PK-index loading and index-before-bloom-filter
    lookup
    order are preserved. This PR now isolates SegmentCache accounting; it
    does not
    include the earlier segment lazy-loading or BF-first lookup
    optimizations.
    
    This limits cache retention, not active-operation memory. Evicted
    Segments remain
    valid through callers' shared pointers, and more frequent eviction can
    increase
    index reloads.
    
    ### Validation
    
    - Added 5 LRU tests and 6 Segment cache tests covering physical
    accounting,
    monotonic/concurrent charge updates, eviction, active handles,
    replacement,
    initialization before/after insertion, cache lifetime and partial-load
    failure.
    - Confirmed the restored rowset loading and lookup code matches baseline
      `695c88b5772`.
    - Changed-file clang-format 16 and build hygiene passed.
    - The unchanged low-level LRU implementation passed standalone
    compilation and
    ASan/UBSan scenarios using external logging/time/value-base stubs. This
    is
      supplementary validation, not a full BE unit-test run.
    - Full BE compilation/unit tests and clang-tidy remain pending CI:
    dependency
    downloads repeatedly timed out before compilation in the local
    environment.
      Test filter: `CacheTest.UpdateCharge*:SegmentPrimaryKeyCacheTest.*`.
    - Previous-build SegmentCache-only cleanup on two idle BEs reduced PK
    cache
      from 9.25/8.93 GiB to a few MiB.
    
    #### Before/after memory validation
    
    The before/after test runs show a change from sustained memory growth to
    repeated memory reclamation:
    
    - **Before:** BE memory continued to grow, and the test eventually
    encountered
    OOM. In the captured dashboard window, the three BEs had mean memory
    usage of
      28.7–35.7 GiB, with the highest observed peak reaching 55.3 GiB.
    - **After:** The test ran stably without OOM. BE memory repeatedly
    dropped after
    transient spikes instead of maintaining the previous upward trend. In
    the
      captured dashboard window, the three BEs had mean memory usage of
      8.33–19.2 GiB, with the highest observed peak reaching 45.3 GiB.
    
    | BE memory metric in the captured window | Before | After |
    | --- | --- | --- |
    | Range of per-BE mean memory usage | 28.7–35.7 GiB | 8.33–19.2 GiB |
    | Highest observed peak across the three BEs | 55.3 GiB | 45.3 GiB |
    | Test behavior | Continued growth, eventually OOM | Repeated memory
    drops, stable operation without OOM |
    
    The dashboard measurements above come from separate runs on different BE
    groups
    and time windows.
    These results validate the observed memory-retention improvement;
    transient
    memory peaks remain, and no throughput comparison is claimed.
    
    ### Release note
    
    Account for pinned PK index roots in SegmentCache eviction to reduce
    persistent
    index memory retention.
---
 be/src/storage/index/primary_key_index.cpp         |  12 +-
 be/src/storage/segment/segment.cpp                 |   6 +-
 be/src/storage/segment/segment.h                   |  16 ++
 be/src/storage/segment/segment_loader.cpp          |  21 ++-
 be/src/util/lru_cache.cpp                          |  47 +++++
 be/src/util/lru_cache.h                            |   7 +
 be/test/storage/cache/lru_cache_test.cpp           | 175 +++++++++++++++++++
 .../segment/segment_primary_key_cache_test.cpp     | 190 +++++++++++++++++++++
 8 files changed, 468 insertions(+), 6 deletions(-)

diff --git a/be/src/storage/index/primary_key_index.cpp 
b/be/src/storage/index/primary_key_index.cpp
index 28fa377e4b4..080e19d503b 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<segment_v2::IndexedColumnReader>(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 af4ec1268a1..0e60f76cc2a 100644
--- a/be/src/storage/segment/segment.cpp
+++ b/be/src/storage/segment/segment.cpp
@@ -742,7 +742,11 @@ Status Segment::load_index(OlapReaderStatistics* stats, 
const io::IOContext* sou
             _pk_index_reader = std::make_unique<PrimaryKeyIndexReader>();
             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 dfdb88ab72c..7a0228b0d29 100644
--- a/be/src/storage/segment/segment.h
+++ b/be/src/storage/segment/segment.h
@@ -22,7 +22,9 @@
 #include <gen_cpp/segment_v2.pb.h>
 #include <glog/logging.h>
 
+#include <atomic>
 #include <cstdint>
+#include <functional>
 #include <map>
 #include <memory> // for unique_ptr
 #include <string>
@@ -214,6 +216,18 @@ public:
     // 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<void(size_t)> callback) {
+        DCHECK(!_cache_charge_callback);
+        _cache_charge_callback = std::move(callback);
+    }
+
     // Variant paths infer their storage type from the physical root reader. 
If an ALTER-added root
     // is absent from an old segment, use the path's declared schema type. 
Other columns use
     // read_type directly.
@@ -322,6 +336,8 @@ private:
     // The memory consumed by querying is tracked in segment iterator.
     int64_t _meta_mem_usage;
     int64_t _tracked_meta_mem_usage = 0;
+    std::atomic<size_t> _pk_index_cache_bytes {0};
+    std::function<void(size_t)> _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 e1c12cbf486..c7f4a1a6d3b 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>(_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<CacheValue*>(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 353c7303212..5d44236c9a5 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<Cache::Handle*>(e);
 }
 
+void LRUCache::update_charge(Cache::Handle* handle, size_t charge) {
+    auto* e = reinterpret_cast<LRUHandle*>(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<LRUHandle*>(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 f01fd0cd6cf..718baabba7f 100644
--- a/be/src/util/lru_cache.h
+++ b/be/src/util/lru_cache.h
@@ -193,6 +193,10 @@ public:
     // 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 @@ public:
                           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 @@ public:
                    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 @@ public:
                    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 a31d532a5a4..5c248838ad6 100644
--- a/be/test/storage/cache/lru_cache_test.cpp
+++ b/be/test/storage/cache/lru_cache_test.cpp
@@ -21,6 +21,7 @@
 #include <gtest/gtest-message.h>
 #include <gtest/gtest-test-part.h>
 
+#include <future>
 #include <iosfwd>
 #include <vector>
 
@@ -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<const CacheValue*>(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<const CacheValue*>(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<CacheValue*>(reinterpret_cast<LRUHandle*>(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<std::future<void>> 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<CacheValueWithKey*>(reinterpret_cast<LRUHandle*>(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 00000000000..7de4aa4625b
--- /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 <gtest/gtest.h>
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+#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<TabletColumnPtr>& columns, 
KeysType keys_type);
+using Generator = std::function<void(size_t rid, int cid, Field& field)>;
+void build_segment(VerticalSegmentWriterOptions opts, TabletSchemaSPtr 
build_schema,
+                   size_t segment_id, TabletSchemaSPtr query_schema, size_t 
nrows,
+                   Generator generator, std::shared_ptr<Segment>* 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<StorageEngine>(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<TYPE_VARCHAR>(std::string(50000, 
'a' + rid));
+            } else {
+                field = 
Field::create_field<TYPE_INT>(static_cast<int32_t>(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_v2::Segment> _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<ErrorCode::KEY_NOT_FOUND>());
+    // 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_v2::Segment>(_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


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to