fix cache_task.h

Author: niushengxiao

src/core/cache_task.h

@@ -15,106 +15,99 @@
 
 namespace cache::task {
 
-enum State { Initial = 0, Working = 1, Finished = 2, Aborted = 3 };
+enum State { Initial, Working, Finished, Aborted };
 
-enum Mode { Write = 1, Read = 2 };
+enum Mode { Write, Read };
 
 class CacheTask;
 
+/**
+ * @brief Represents a single block of data within a cache task
+ * 
+ * Each CacheBlock corresponds to a fixed-size chunk of KV cache data that can be
+ * independently read from or written to storage. Blocks are processed asynchronously
+ * by worker threads and track their own state throughout the operation lifecycle.
+ */
 class CacheBlock {
 public:
-  /// @brief Constructor, calculates the SHA-256 hash of the data.
-  /// @param hash Data hash.
-  /// @param m Mode: Read or Write
-  /// @param task Corresponding CacheTask
   CacheBlock(std::string hash_, const int64_t block_idx_, CacheTask *task_)
       : hash(std::move(hash_)), task(task_), block_idx(block_idx_) {}
 
-  /// @brief 禁止拷贝和移动
-  CacheBlock(CacheBlock &&other) = delete;
-  CacheBlock &operator=(CacheBlock &&other) = delete;
-
   bool ready() const { return state == State::Finished; }
 
   int64_t block_idx;
-  CacheTask *task;  ///< Corresponding Task
-  std::string hash; ///< Hash of the block.
-  State state{};    ///< Read/write state of the block.
+  CacheTask *task;
+  std::string hash;
+  State state{};
 };
 
-/// @brief Task cache class, storing only cache blocks.
+/**
+ * @brief Manages a collection of cache blocks for a single read or write operation
+ * 
+ * CacheTask represents a complete cache operation request from Python code, containing
+ * multiple blocks that are processed in parallel by worker threads. It tracks completion
+ * status, provides thread-safe access to shared state, and manages the lifecycle of all
+ * associated blocks through RAII-managed unique_ptr ownership.
+ * 
+ * Thread Safety: Atomic members (num_finished_blocks, num_data_ready_blocks, completion_notified)
+ * are lock-free for high-frequency access. The state_mutex protects page_already_list_ updates.
+ */
 class CacheTask {
 public:
-  CacheTask() = delete;
-
-  /// @brief Constructor, determines mode from user input ('r' or 'w').
-  /// @param hashs Hash sequence.
-  /// @param mode_str Mode string: "r" for Read, "w" for Write
   CacheTask(const std::vector<std::string> &hashs, torch::Tensor kv_page_indexer, const std::string &mode_str)
-      : num_finished_blocks(0), num_data_ready_blocks(0), page_indexer(std::move(kv_page_indexer)),
-        completion_notified(false) {
+      : num_finished_blocks(0)
+      , num_data_ready_blocks(0)
+      , page_indexer(std::move(kv_page_indexer))
+      , completion_notified(false) {
 
     if (mode_str == "r") {
-      mode = Mode::Read;
+      operation_mode = Mode::Read;
     } else if (mode_str == "w") {
-      mode = Mode::Write;
+      operation_mode = Mode::Write;
     } else {
       throw std::invalid_argument("Invalid mode string. Use 'r' for Read or 'w' for Write.");
     }
 
     blocks.reserve(hashs.size());
-    for (int64_t idx = 0; idx < hashs.size(); ++idx) {
-      blocks.emplace_back(new CacheBlock(hashs[idx], idx, this));
+    int64_t idx = 0;
+    for (const auto& hash : hashs) {
+      blocks.emplace_back(std::make_unique<CacheBlock>(hash, idx++, this));
     }
   }
 
-  ~CacheTask() {
-    for (auto block : blocks) {
-      delete block;
-    }
-  }
-
-  /// @brief 禁止拷贝和移动
-  CacheTask(CacheTask &&other) = delete;
-  CacheTask &operator=(CacheTask &&other) = delete;
-
   bool ready() const { return num_finished_blocks.load(std::memory_order_acquire) == blocks.size(); }
 
-  /// @brief Check if data is safe to release pages (for write mode)
-  /// For write mode: returns true when data has been copied from KV cache
-  /// For read mode: equivalent to ready()
   bool data_safe() const {
-    if (mode == Mode::Write) {
+    if (operation_mode == Mode::Write) {
       return num_data_ready_blocks.load(std::memory_order_acquire) >= static_cast<int64_t>(blocks.size());
     }
-    return ready(); // For read mode, data_safe is same as ready
+    return ready();
   }
 
   bool mark_completion_notified() { return !completion_notified.exchange(true, std::memory_order_acq_rel); }
 
-  std::vector<State> state() {
-    auto ret = std::vector<State>(blocks.size());
-    for (int32_t i = 0; i < blocks.size(); ++i) {
-      ret[i] = blocks[i]->state;
+  std::vector<State> state() const {
+    std::vector<State> ret;
+    ret.reserve(blocks.size());
+    for (const auto& block : blocks) {
+      ret.push_back(block->state);
     }
     return ret;
   }
 
   std::vector<int32_t> get_page_already_list() const {
-    std::lock_guard<std::mutex> lock_guard(const_cast<std::mutex &>(lock));
+    std::lock_guard<std::mutex> lock_guard(state_mutex);
     return page_already_list;
   }
 
-  // 这个指针是用来标记 task 中的数据存取位置的
   torch::Tensor page_indexer;
-
-  std::mutex lock;                            ///< Task state lock
-  std::vector<CacheBlock *> blocks;           ///< Blocks stored as shared_ptr
-  std::atomic<int64_t> num_finished_blocks;   ///< Number of finished blocks (atomic for thread-safe reading)
-  std::atomic<int64_t> num_data_ready_blocks; ///< Number of blocks with data copied (for write mode)
-  Mode mode;                                  ///< Read/write mode of the task.
+  mutable std::mutex state_mutex;
+  std::vector<std::unique_ptr<CacheBlock>> blocks;
+  std::atomic<int64_t> num_finished_blocks;
+  std::atomic<int64_t> num_data_ready_blocks;
+  Mode operation_mode;
   std::atomic<bool> completion_notified;
-  std::vector<int32_t> page_already_list; ///< List of page indices already persisted to disk (for write mode)
+  std::vector<int32_t> page_already_list;
 };
 
 } // namespace cache::task

src/core/task_queue.h

@@ -25,8 +25,8 @@ class TaskQueue {
   TaskQueue &operator=(TaskQueue &&other) = delete;
 
   cache::error::LMError_t submit(const std::shared_ptr<cache::task::CacheTask> &task) {
-    for (cache::task::CacheBlock *block : task->blocks) {
-      blocks_.push(block);
+    for (const auto &block : task->blocks) {
+      blocks_.push(block.get());
     }
     return cache::error::LM_SUCCESS;
   }

src/service/cache_service.h

@@ -269,16 +269,16 @@ class CacheService {
    * read/write operations on the kvcache.
    */
   void abort_task(const std::shared_ptr<cache::task::CacheTask> &task) {
-    for (cache::task::CacheBlock *block : task->blocks) {
-      abort(block);
+    for (const auto &block : task->blocks) {
+      abort(block.get());
     }
   }
 
   /**
    * Notify the system to immediately abandon the subsequent execution of a Block.
    */
   void abort(cache::task::CacheBlock *block) {
-    std::lock_guard<std::mutex> lock(block->task->lock);
+    std::lock_guard<std::mutex> lock(block->task->state_mutex);
     if (block->state == cache::task::State::Initial || block->state == cache::task::State::Working) {
       block->state = cache::task::State::Aborted;
       block->task->num_finished_blocks.fetch_add(1, std::memory_order_release);
@@ -291,7 +291,7 @@ class CacheService {
   void deliver(cache::task::CacheBlock *block) {
     auto task = block->task;
     {
-      std::lock_guard<std::mutex> lock(task->lock);
+      std::lock_guard<std::mutex> lock(task->state_mutex);
 
       if (block->state == cache::task::State::Working) {
         block->state = cache::task::State::Finished;
@@ -343,15 +343,23 @@ inline void CacheService::finalize_task(cache::task::CacheTask *task) {
     std::lock_guard<std::mutex> lock(lock_);
     for (auto it = taskpool_.begin(); it != taskpool_.end(); ++it) {
       if (it->get() == task) {
+        // IMPORTANT: After this erase, the shared_ptr may be destroyed if Python side has released it.
+        // This is safe because:
+        // 1. All blocks have finished (checked by task->ready())
+        // 2. No worker threads should be accessing this task anymore
+        // 3. The task pointer 'task' is only used for comparison and counter update
         taskpool_.erase(it);
         break;
       }
     }
   }
 
   std::atomic<int64_t> *active_counter =
-      (task->mode == cache::task::Mode::Read) ? &active_read_creates_ : &active_write_creates_;
+      (task->operation_mode == cache::task::Mode::Read) ? &active_read_creates_ : &active_write_creates_;
   active_counter->fetch_sub(1, std::memory_order_relaxed);
+  
+  // Note: 'task' pointer may become invalid after this point if no other references exist
+  // on_task_finalized should not dereference 'task' beyond this point unless it maintains its own reference
   on_task_finalized(task);
 }
 

src/service/local_cache_service.h

@@ -62,12 +62,18 @@ class LocalCacheService : public CacheService {
 
     storage_ = make_unique<LocalStorageEngine>(file, storage_size, num_shard, block_size_);
 
-    r_cpu_buffers_.resize(num_workers_);
-    w_cpu_buffers_.resize(num_workers_);
-
-    for (size_t i = 0; i < num_workers_; ++i) {
-      r_cpu_buffers_[i] = new char[block_size_];
-      w_cpu_buffers_[i] = new char[block_size_];
+    // Use unique_ptr for exception safety - if any allocation fails, previous allocations are automatically cleaned up
+    r_cpu_buffers_.reserve(num_workers_);
+    w_cpu_buffers_.reserve(num_workers_);
+
+    try {
+      for (size_t i = 0; i < num_workers_; ++i) {
+        r_cpu_buffers_.emplace_back(new char[block_size_]);
+        w_cpu_buffers_.emplace_back(new char[block_size_]);
+      }
+    } catch (...) {
+      // unique_ptr will automatically clean up already allocated buffers
+      throw;
     }
   }
 
@@ -82,13 +88,7 @@ class LocalCacheService : public CacheService {
       }
     }
 
-    for (auto &buffer : r_cpu_buffers_) {
-      delete[] buffer;
-    }
-    for (auto &buffer : w_cpu_buffers_) {
-      delete[] buffer;
-    }
-
+    // unique_ptr will automatically delete the buffers
     r_cpu_buffers_.clear();
     w_cpu_buffers_.clear();
   }
@@ -120,7 +120,7 @@ class LocalCacheService : public CacheService {
 
 protected:
   void on_task_finalized(cache::task::CacheTask *task) override {
-    if (task->mode == cache::task::Mode::Write) {
+    if (task->operation_mode == cache::task::Mode::Write) {
       // Try to acquire the lock, skip logging if contention occurs
       std::unique_lock<std::mutex> guard(log_mutex_, std::try_to_lock);
       if (!guard.owns_lock()) {
@@ -174,7 +174,7 @@ class LocalCacheService : public CacheService {
       return;
     }
 
-    if (task->mode != cache::task::Mode::Read) {
+    if (task->operation_mode != cache::task::Mode::Read) {
       return;
     }
     if (active_read_creates_.load(std::memory_order_relaxed) != 0) {
@@ -239,7 +239,7 @@ class LocalCacheService : public CacheService {
       if (auto block = this->queue_->claim()) {
         if (block != nullptr) {
           CacheTask *task = block->task;
-          char *cpu_buffer = (task->mode == Mode::Read) ? r_cpu_buffers_[index] : w_cpu_buffers_[index];
+          char *cpu_buffer = (task->operation_mode == Mode::Read) ? r_cpu_buffers_[index].get() : w_cpu_buffers_[index].get();
           processTask(block, cpu_buffer);
         }
       }
@@ -266,7 +266,7 @@ class LocalCacheService : public CacheService {
     }
 
     bool success = false;
-    if (task->mode == Mode::Read) {
+    if (task->operation_mode == Mode::Read) {
       success = handleReadCpu(block, cpu_buffer, page_ptr, num_of_page);
     } else {
       success = handleWriteCpu(block, cpu_buffer, page_ptr, num_of_page);
@@ -291,7 +291,7 @@ class LocalCacheService : public CacheService {
     }
 
     {
-      std::lock_guard<std::mutex> lock(block->task->lock);
+      std::lock_guard<std::mutex> lock(block->task->state_mutex);
       if (block->state != cache::task::State::Working) {
         return false;
       }
@@ -311,7 +311,7 @@ class LocalCacheService : public CacheService {
 
   bool handleWriteCpu(CacheBlock *block, char *cpu_buffer, int32_t *page_ptr, int64_t num_of_page) {
     {
-      std::lock_guard<std::mutex> lock(block->task->lock);
+      std::lock_guard<std::mutex> lock(block->task->state_mutex);
       if (block->state != cache::task::State::Working) {
         return false;
       }
@@ -395,8 +395,8 @@ class LocalCacheService : public CacheService {
   vector<thread> workers_;                              ///< Worker threads
   bool stop_;                                           ///< Thread stop flag
   size_t num_workers_;                                  ///< Number of worker threads
-  vector<char *> r_cpu_buffers_;                        ///< CPU buffers for read worker
-  vector<char *> w_cpu_buffers_;                        ///< CPU buffers for write worker
+  vector<unique_ptr<char[]>> r_cpu_buffers_;            ///< CPU buffers for read worker (RAII managed)
+  vector<unique_ptr<char[]>> w_cpu_buffers_;            ///< CPU buffers for write worker (RAII managed)
   std::atomic<uint64_t> total_written_bytes_;           ///< Total bytes written to disk
   std::atomic<int64_t> first_write_time_ticks_;         ///< First write start time in steady clock ticks
   std::atomic<int64_t> last_write_time_ticks_;          ///< Last write completion time in steady clock ticks

src/storage/local_storage_engine.h

@@ -164,23 +164,21 @@ class LocalStorageEngine : public StorageEngine {
     files_.resize(shard_);
     file_fds_.resize(shard_, -1);
 
-    for (size_t i = 0; i < shard_; i++) {
-      caches_[i] = std::make_shared<LocalCacheIndex>(shard_capacity);
-      locks_[i] = std::make_shared<std::mutex>();
+    try {
+      for (size_t i = 0; i < shard_; i++) {
+        caches_[i] = std::make_shared<LocalCacheIndex>(shard_capacity);
+        locks_[i] = std::make_shared<std::mutex>();
+      }
+      createOrOpenFiles(shard_storage_size);
+    } catch (...) {
+      // Clean up any partially opened files on exception
+      cleanup();
+      throw;
     }
-    createOrOpenFiles(shard_storage_size);
   }
 
   ~LocalStorageEngine() override {
-    for (size_t i = 0; i < shard_; i++) {
-      if (files_[i].is_open()) {
-        files_[i].close();
-      }
-      if (file_fds_[i] >= 0) {
-        close(file_fds_[i]);
-        file_fds_[i] = -1;
-      }
-    }
+    cleanup();
   }
 
   bool query(const std::string &hash) override {
@@ -251,6 +249,23 @@ class LocalStorageEngine : public StorageEngine {
 private:
   inline size_t getShard(const std::string &hash) { return std::hash<std::string>{}(hash) % shard_; }
 
+  // Helper function to clean up file resources
+  void cleanup() {
+    for (size_t i = 0; i < shard_; i++) {
+      if (files_[i].is_open()) {
+        try {
+          files_[i].close();
+        } catch (...) {
+          // Ignore exceptions during cleanup
+        }
+      }
+      if (file_fds_[i] >= 0) {
+        close(file_fds_[i]);
+        file_fds_[i] = -1;
+      }
+    }
+  }
+
   void createOrOpenFiles(size_t shard_storage_size) {
     for (size_t i = 0; i < shard_; i++) {
       std::stringstream ss;