[opt] Prevent Load starvation by reserving buffers for high-priority operations (#895)

## Purpose
This PR introduces a reserved buffer mechanism to prevent Load
operations from being starved when Dump operations consume all available
buffers slowly.
In scenarios where Dump operations process buffers slowly, they can
occupy all buffers, causing Load operations to wait indefinitely. By
reserving a portion of buffers exclusively for Load operations, we
ensure system responsiveness and prevent deadlocks.

## Modifications
- Added loadExclusiveBufferNumber configuration parameter (default 1024)
to specify the number of buffers reserved for Load operations.
- Modified buffer allocation logic to respect reserved buffers: ordinary
allocations cannot use the last loadExclusiveBufferNumber buffers, while
Load operations (with allowReserved=true) can access all buffers.
- Updated LoadQueue::DispatchOneTask to request buffers with
allowReserved=true to ensure Load tasks can always get buffers.
- Added comprehensive tests for reserved buffer behavior.

## Test
- Added GetReservedNode test to verify reserved buffer allocation logic.
- Existing tests pass with updated configurations.

## Configuration
New config parameter: `cache_load_exclusive_buffer_number` (default:
1024)

Author: mag1c-h

ucm/store/cache/cc/cache_store.cc

@@ -128,6 +128,7 @@ class CacheStore : public StoreV1 {
         config.GetNumber("running_queue_depth", param.runningQueueDepth);
         config.GetNumber("timeout_ms", param.timeoutMs);
         config.GetNumber("cache_stream_number", param.streamNumber);
+        config.GetNumber("cache_load_exclusive_buffer_number", param.loadExclusiveBufferNumber);
         return param;
     }
     Status CheckSizeConfig(const Config& config)
@@ -160,7 +161,7 @@ class CacheStore : public StoreV1 {
         auto s = CheckSizeConfig(config);
         if (s.Failure()) { return s; }
         auto bufferNumber = config.bufferCapacity / config.shardSize;
-        if (bufferNumber < 1024) {
+        if (bufferNumber < 1024 || bufferNumber < config.loadExclusiveBufferNumber * 2) {
             return Status::InvalidParam("too small buffer({}) on shard({})", config.bufferCapacity,
                                         config.shardSize);
         }

ucm/store/cache/cc/global_config.h

@@ -39,6 +39,7 @@ struct Config {
     bool ioDirect{false};
     std::vector<ssize_t> cpuAffinityCores{};
     size_t bufferCapacity{256ULL << 30};
+    size_t loadExclusiveBufferNumber{1024};
     bool shareBufferEnable{true};
     size_t waitingQueueDepth{8192};
     size_t runningQueueDepth{524288};

ucm/store/cache/cc/load_queue.cc

@@ -86,7 +86,7 @@ void LoadQueue::DispatchOneTask(TaskPair&& pair)
     for (size_t i = 0; i < nShard; i++) {
         auto& shard = task->desc[i];
         ShardTask shardTask;
-        shardTask.bufferHandle = buffer_->Get(shard.owner, shard.index);
+        shardTask.bufferHandle = buffer_->Get(shard.owner, shard.index, true);
         shardTask.backendTaskHandle = 0;
         if (shardTask.bufferHandle.Owner() && !shardTask.bufferHandle.Ready()) {
             Detail::TaskDesc backendTask{

ucm/store/cache/cc/trans_buffer.cc

@@ -70,12 +70,13 @@ class BufferStrategy {
         int32_t deviceId{-1};
         size_t nodeSize{0};
         size_t totalSize{0};
+        size_t reservedNumber{0};
     };
     BaseConfig base_;
 
 public:
-    BufferStrategy(int32_t deviceId, size_t nodeSize, size_t totalSize)
-        : base_({deviceId, nodeSize, totalSize})
+    BufferStrategy(int32_t deviceId, size_t nodeSize, size_t totalSize, size_t reservedNumber)
+        : base_({deviceId, nodeSize, totalSize, reservedNumber})
     {
     }
     virtual ~BufferStrategy() = default;
@@ -86,7 +87,7 @@ class BufferStrategy {
     virtual void NodeLock(size_t iNode) = 0;
     virtual void NodeUnlock(size_t iNode) = 0;
     virtual size_t& FirstAt(size_t iBucket) = 0;
-    virtual size_t FetchNode() = 0;
+    virtual size_t FetchNode(bool allowReserved) = 0;
     virtual void* DataAt(size_t iNode) = 0;
     virtual BufferMetaNode* MetaAt(size_t iNode) = 0;
 };
@@ -132,8 +133,9 @@ class LocalBufferStrategy : public BufferStrategy {
     std::shared_ptr<void> data_;
 
 public:
-    LocalBufferStrategy(int32_t deviceId, size_t nodeSize, size_t totalSize, bool ioDirect)
-        : BufferStrategy(deviceId, nodeSize, totalSize), ioDirect_(ioDirect)
+    LocalBufferStrategy(int32_t deviceId, size_t nodeSize, size_t totalSize, size_t reservedNumber,
+                        bool ioDirect)
+        : BufferStrategy(deviceId, nodeSize, totalSize, reservedNumber), ioDirect_(ioDirect)
     {
     }
     Status Setup() override
@@ -181,10 +183,11 @@ class LocalBufferStrategy : public BufferStrategy {
     void NodeLock(size_t iNode) override { nodeLocks_[iNode].Lock(); }
     void NodeUnlock(size_t iNode) override { nodeLocks_[iNode].Unlock(); }
     size_t& FirstAt(size_t iBucket) override { return header_.buckets[iBucket]; }
-    size_t FetchNode() override
+    size_t FetchNode(bool allowReserved) override
     {
-        auto head = header_.freeHead++;
-        if (header_.freeHead == header_.nNode) { header_.freeHead = 0; }
+        const auto limit = header_.nNode - (allowReserved ? 0 : base_.reservedNumber);
+        const auto head = header_.freeHead;
+        header_.freeHead = (head + 1 == limit) ? 0 : (head + 1);
         return head;
     }
     void* DataAt(size_t iNode) override
@@ -368,8 +371,8 @@ class SharedBufferStrategy : public BufferStrategy {
 
 public:
     SharedBufferStrategy(const std::string& uuid, int32_t deviceId, size_t nodeSize,
-                         size_t totalSize)
-        : BufferStrategy(deviceId, nodeSize, totalSize), uuid_(uuid)
+                         size_t totalSize, size_t reservedNumber)
+        : BufferStrategy(deviceId, nodeSize, totalSize, reservedNumber), uuid_(uuid)
     {
     }
     ~SharedBufferStrategy() override
@@ -410,11 +413,12 @@ class SharedBufferStrategy : public BufferStrategy {
     void NodeLock(size_t iNode) override { header_->nodeLocks[iNode].Lock(); }
     void NodeUnlock(size_t iNode) override { header_->nodeLocks[iNode].Unlock(); }
     size_t& FirstAt(size_t iBucket) override { return header_->buckets[iBucket]; }
-    size_t FetchNode() override
+    size_t FetchNode(bool allowReserved) override
     {
+        const auto limit = header_->nNode - (allowReserved ? 0 : base_.reservedNumber);
         header_->lock.Lock();
-        auto iNode = header_->freeHead++;
-        if (header_->freeHead == nNode_) { header_->freeHead = 0; }
+        const auto iNode = header_->freeHead;
+        header_->freeHead = (iNode + 1 == limit) ? 0 : (iNode + 1);
         header_->lock.Unlock();
         return iNode;
     }
@@ -424,7 +428,9 @@ class SharedBufferStrategy : public BufferStrategy {
 
 class SharedBufferWatcherStrategy : public SharedBufferStrategy {
 public:
-    SharedBufferWatcherStrategy(const std::string& uuid) : SharedBufferStrategy(uuid, -1, 0, 0) {}
+    SharedBufferWatcherStrategy(const std::string& uuid) : SharedBufferStrategy(uuid, -1, 0, 0, 0)
+    {
+    }
     Status Setup() override
     {
         shmName_ = ShmPrefix() + uuid_;
@@ -462,10 +468,12 @@ Status TransBuffer::Setup(const Config& config)
     try {
         if (!config.shareBufferEnable) {
             strategy_ = std::make_shared<LocalBufferStrategy>(
-                config.deviceId, config.shardSize, config.bufferCapacity, config.ioDirect);
+                config.deviceId, config.shardSize, config.bufferCapacity,
+                config.loadExclusiveBufferNumber, config.ioDirect);
         } else if (config.deviceId >= 0) {
             strategy_ = std::make_shared<SharedBufferStrategy>(
-                config.uniqueId, config.deviceId, config.shardSize, config.bufferCapacity);
+                config.uniqueId, config.deviceId, config.shardSize, config.bufferCapacity,
+                config.loadExclusiveBufferNumber);
         } else {
             strategy_ = std::make_shared<SharedBufferWatcherStrategy>(config.uniqueId);
         }
@@ -475,7 +483,8 @@ Status TransBuffer::Setup(const Config& config)
     return strategy_->Setup();
 }
 
-TransBuffer::Handle TransBuffer::Get(const Detail::BlockId& blockId, size_t shardIdx)
+TransBuffer::Handle TransBuffer::Get(const Detail::BlockId& blockId, size_t shardIdx,
+                                     bool allowReserved)
 {
     auto iBucket = Hash(blockId, shardIdx);
     bool owner = false;
@@ -485,7 +494,7 @@ TransBuffer::Handle TransBuffer::Get(const Detail::BlockId& blockId, size_t shar
         strategy_->BucketUnlock(iBucket);
         return Handle{this, iNode, owner};
     }
-    iNode = Alloc(blockId, shardIdx, iBucket);
+    iNode = Alloc(blockId, shardIdx, iBucket, allowReserved);
     strategy_->BucketUnlock(iBucket);
     return Handle(this, iNode, true);
 }
@@ -536,10 +545,11 @@ size_t TransBuffer::FindAt(size_t iBucket, const Detail::BlockId& blockId, size_
     return iNode;
 }
 
-size_t TransBuffer::Alloc(const Detail::BlockId& blockId, size_t shardIdx, size_t iBucket)
+size_t TransBuffer::Alloc(const Detail::BlockId& blockId, size_t shardIdx, size_t iBucket,
+                          bool allowReserved)
 {
     for (;;) {
-        auto iNode = strategy_->FetchNode();
+        auto iNode = strategy_->FetchNode(allowReserved);
         auto meta = strategy_->MetaAt(iNode);
         strategy_->NodeLock(iNode);
         if (meta->reference > 0) {

ucm/store/cache/cc/trans_buffer.h

@@ -98,13 +98,14 @@ class TransBuffer {
 
 public:
     Status Setup(const Config& config);
-    Handle Get(const Detail::BlockId& blockId, size_t shardIdx);
+    Handle Get(const Detail::BlockId& blockId, size_t shardIdx, bool allowReserved = false);
     bool Exist(const Detail::BlockId& blockId, size_t shardIdx);
 
 private:
     bool ExistAt(size_t iBucket, const Detail::BlockId& blockId, size_t shardIdx);
     size_t FindAt(size_t iBucket, const Detail::BlockId& blockId, size_t shardIdx, bool& owner);
-    size_t Alloc(const Detail::BlockId& blockId, size_t shardIdx, size_t iBucket);
+    size_t Alloc(const Detail::BlockId& blockId, size_t shardIdx, size_t iBucket,
+                 bool allowReserved = false);
     void MoveTo(size_t iBucket, size_t iNode);
     void Remove(size_t iBucket, size_t iNode);
     void* DataAt(Index pos);

ucm/store/test/case/cache/cache_trans_buffer_test.cc

@@ -42,6 +42,7 @@ TEST_P(UCCacheTransBufferTest, GetFirstNode)
     config.bufferCapacity = config.shardSize * 32768;
     config.shareBufferEnable = GetParam();
     config.deviceId = 0;
+    config.loadExclusiveBufferNumber = 0;
     auto s = transBuffer.Setup(config);
     ASSERT_EQ(s, UC::Status::OK());
     auto blockId = UC::Test::Detail::TypesHelper::MakeBlockId("a1b2c3d4e5f6789012345678901234ab");
@@ -59,6 +60,44 @@ TEST_P(UCCacheTransBufferTest, GetFirstNode)
     ASSERT_TRUE(handle2.Ready());
 }
 
+TEST_P(UCCacheTransBufferTest, GetReservedNode)
+{
+    UC::CacheStore::TransBuffer transBuffer;
+    UC::CacheStore::Config config;
+    config.uniqueId = rd.RandomString(10);
+    config.shardSize = 32768;
+    config.loadExclusiveBufferNumber = 16;
+    config.bufferCapacity = config.shardSize * (config.loadExclusiveBufferNumber + 1);
+    config.shareBufferEnable = GetParam();
+    config.deviceId = 0;
+    auto s = transBuffer.Setup(config);
+    ASSERT_EQ(s, UC::Status::OK());
+    auto blockId1 = UC::Test::Detail::TypesHelper::MakeBlockId("a1b2c3d4e5f6789012345678901234ab");
+    auto blockId2 = UC::Test::Detail::TypesHelper::MakeBlockId("a2b2c3d4e5f6789012345678901234ab");
+    constexpr size_t shardIdx = 0;
+    void* ptr = nullptr;
+    {
+        auto handle1 = transBuffer.Get(blockId1, shardIdx);
+        ASSERT_TRUE(handle1);
+        ptr = handle1.Data();
+    }
+    {
+        auto handle2 = transBuffer.Get(blockId2, shardIdx);
+        ASSERT_TRUE(handle2);
+        ASSERT_EQ(ptr, handle2.Data());
+    }
+    {
+        auto handle1 = transBuffer.Get(blockId1, shardIdx, true);
+        ASSERT_TRUE(handle1);
+        ptr = handle1.Data();
+    }
+    {
+        auto handle2 = transBuffer.Get(blockId2, shardIdx, true);
+        ASSERT_TRUE(handle2);
+        ASSERT_NE(ptr, handle2.Data());
+    }
+}
+
 TEST_P(UCCacheTransBufferTest, InsertDifferentDataRepeatedly)
 {
     constexpr size_t nBatch = 2;
@@ -71,6 +110,7 @@ TEST_P(UCCacheTransBufferTest, InsertDifferentDataRepeatedly)
     config.bufferCapacity = nBlock * nShard * config.shardSize;
     config.shareBufferEnable = GetParam();
     config.deviceId = 0;
+    config.loadExclusiveBufferNumber = 0;
     auto s = transBuffer.Setup(config);
     ASSERT_EQ(s, UC::Status::OK());
     for (size_t iBatch = 0; iBatch < nBatch; iBatch++) {