[fix](mow) Account for pinned PK roots in segment cache eviction

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();
 }

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<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

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 <optional>
@@ -181,6 +183,18 @@ class Segment : public std::enable_shared_from_this<Segment>, 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<void(size_t)> callback) {
+        DCHECK(!_cache_charge_callback);
+        _cache_charge_callback = std::move(callback);
+    }
+
     // Variant paths use segment metadata; other columns use `read_type`.
     std::shared_ptr<const IDataType> get_data_type_of(const TabletColumn& read_column,
                                                       const DataTypePtr& read_type,
@@ -292,6 +306,8 @@ class Segment : public std::enable_shared_from_this<Segment>, 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<size_t> _pk_index_cache_bytes {0};
+    std::function<void(size_t)> _cache_charge_callback;
 
     RowsetId _rowset_id;
     TabletSchemaSPtr _tablet_schema;

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);
 }

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);

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; };

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);