Optimize ASU task manager with lock-free slots

Author: AooooooA-C

ucm/transport/kv/asu/common/task_manager_base.h

@@ -24,20 +24,43 @@
 #pragma once
 
 #include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
 #include <memory>
-#include <mutex>
 #include <string>
-#include <unordered_map>
 #include <utility>
+#include <vector>
+#include <functional>
+
 #include "asu_transport/types.h"
 
 namespace UC::ASU {
 
 template <typename Context, typename State>
 class TaskManagerBase {
 public:
-    TaskManagerBase(State initialState, std::string taskName)
-        : initialState_(initialState), taskName_(std::move(taskName))
+    static constexpr std::size_t kMinSlotCount = 1024;
+    static constexpr std::size_t kDefaultSlotCount = 8192;
+
+    static std::size_t RecommendSlotCount(std::size_t maxInflightTasks)
+    {
+        // Keep load factor <= 0.5 for open addressing.
+        // For example: 4096 inflight tasks -> 8192 slots.
+        const auto required = std::max<std::size_t>(
+            kMinSlotCount,
+            maxInflightTasks * 2);
+        return NormalizeSlotCount(required);
+    }
+
+    explicit TaskManagerBase(
+        State initialState,
+        std::string taskName,
+        std::size_t slotCount = kDefaultSlotCount)
+        : initialState_(initialState),
+          taskName_(std::move(taskName)),
+          slots_(NormalizeSlotCount(slotCount)),
+          slotMask_(slots_.size() - 1)  // Used for efficient slot index calculation: bitwise modulo
     {
     }
 
@@ -51,27 +74,89 @@ class TaskManagerBase {
         auto sharedCtx = std::shared_ptr<Context>(std::move(ctx));
         sharedCtx->state.store(initialState_, std::memory_order_release);
 
-        std::lock_guard<std::mutex> lock(mutex_);
+        TaskId newTaskId = kInvalidTaskId;
         do {
-            taskId = nextTaskId_.fetch_add(1, std::memory_order_relaxed);
-        } while (taskId == kInvalidTaskId || tasks_.find(taskId) != tasks_.end());
+            newTaskId = nextTaskId_.fetch_add(1, std::memory_order_relaxed);
+        } while (newTaskId == kInvalidTaskId);  // kInvalidTaskId is 0, so task id starts from 1 to avoid allocating invalid IDs
+
+        sharedCtx->taskId = newTaskId;
+
+        const auto start = Hash(newTaskId) & slotMask_;
+        const auto capacity = slots_.size();
+
+        for (std::size_t probe = 0; probe < capacity; ++probe) {
+            auto& slot = slots_[(start + probe) & slotMask_];
+
+            // CAS: Try to transition EMPTY → WRITING
+            std::uint8_t expected = SlotState::EMPTY;
+            if (!slot.state.compare_exchange_strong(
+                    expected,
+                    SlotState::WRITING,
+                    std::memory_order_acq_rel,
+                    std::memory_order_acquire)) {
+                continue;
+            }
+
+            AtomicStoreCtx(slot, sharedCtx, std::memory_order_release);
+            slot.taskId.store(newTaskId, std::memory_order_release);
+            slot.state.store(SlotState::READY, std::memory_order_release);
+
+            taskId = newTaskId;
+            return Status::OK();
+        }
+
+        taskId = kInvalidTaskId;
 
-        sharedCtx->taskId = taskId;
-        tasks_.emplace(taskId, std::move(sharedCtx));
-        return Status::OK();
+        // Consider adding RESOURCE_EXHAUSTED / NO_SPACE error codes to StatusCode
+        return Status::Error(
+            StatusCode::INVALID_ARGUMENT,
+            taskName_ + " task table is full");
     }
 
     std::shared_ptr<Context> Get(TaskId taskId)
     {
-        std::lock_guard<std::mutex> lock(mutex_);
-        auto iter = tasks_.find(taskId);
-        if (iter == tasks_.end()) { return nullptr; }
-        return iter->second;
+        if (taskId == kInvalidTaskId) {
+            return nullptr;
+        }
+
+        const auto start = Hash(taskId) & slotMask_;
+        const auto capacity = slots_.size();
+
+        for (std::size_t probe = 0; probe < capacity; ++probe) {
+            auto& slot = slots_[(start + probe) & slotMask_];
+
+            const auto state1 = slot.state.load(std::memory_order_acquire);
+            if (state1 != SlotState::READY) {
+                continue;
+            }
+
+            const auto id1 = slot.taskId.load(std::memory_order_acquire);
+            if (id1 != taskId) {
+                continue;
+            }
+
+            auto ptr = AtomicLoadCtx(slot, std::memory_order_acquire);
+            if (!ptr) {
+                continue;
+            }
+
+            // Double-check to avoid returning a ctx from a reused slot.
+            const auto id2 = slot.taskId.load(std::memory_order_acquire);
+            const auto state2 = slot.state.load(std::memory_order_acquire);
+
+            if (state2 == SlotState::READY &&
+                id2 == taskId &&
+                ptr->taskId == taskId) {
+                return ptr;
+            }
+        }
+
+        return nullptr;
     }
 
     std::vector<std::shared_ptr<Context>> GetAll()
     {
-        std::lock_guard<std::mutex> lock(mutex_);
+        std::lock_guard<std::mutex> lock(mu_);
         std::vector<std::shared_ptr<Context>> tasks;
         tasks.reserve(tasks_.size());
         for (const auto& item : tasks_) { tasks.emplace_back(item.second); }
@@ -80,21 +165,114 @@ class TaskManagerBase {
 
     Status Remove(TaskId taskId)
     {
-        std::lock_guard<std::mutex> lock(mutex_);
-        auto erased = tasks_.erase(taskId);
-        if (erased == 0) {
-            return Status::Error(StatusCode::TASK_NOT_FOUND, taskName_ + " task not found");
+        if (taskId == kInvalidTaskId) {
+            return Status::Error(
+                StatusCode::TASK_NOT_FOUND,
+                taskName_ + " task not found");
+        }
+
+        const auto start = Hash(taskId) & slotMask_;
+        const auto capacity = slots_.size();
+
+        for (std::size_t probe = 0; probe < capacity; ++probe) {
+            auto& slot = slots_[(start + probe) & slotMask_];
+
+            const auto state = slot.state.load(std::memory_order_acquire);
+            if (state != SlotState::READY) {
+                continue;  // Only process slots in READY state
+            }
+
+            const auto id = slot.taskId.load(std::memory_order_acquire);
+            if (id != taskId) {
+                continue;
+            }
+
+            std::uint8_t expected = SlotState::READY;
+            if (!slot.state.compare_exchange_strong(
+                    expected,
+                    SlotState::REMOVING,
+                    std::memory_order_acq_rel,
+                    std::memory_order_acquire)) {
+                continue;  // CAS failed, continue probing
+            }
+
+            AtomicStoreCtx(slot, std::shared_ptr<Context>{}, std::memory_order_release);
+            slot.taskId.store(kInvalidTaskId, std::memory_order_release);
+            slot.state.store(SlotState::EMPTY, std::memory_order_release);
+
+            return Status::OK();
+        }
+
+        return Status::Error(
+            StatusCode::TASK_NOT_FOUND,
+            taskName_ + " task not found");
+    }
+
+private:
+    struct SlotState {
+        static constexpr std::uint8_t EMPTY = 0;
+        static constexpr std::uint8_t WRITING = 1;
+        static constexpr std::uint8_t READY = 2;
+        static constexpr std::uint8_t REMOVING = 3;
+    };
+    // Ensure each Slot is aligned to 64 bytes to avoid False Sharing
+    struct alignas(64) Slot {
+        std::atomic<std::uint8_t> state{SlotState::EMPTY};
+        std::atomic<TaskId> taskId{kInvalidTaskId};
+
+        // Use atomic_load/atomic_store free functions for shared_ptr.
+        // This avoids requiring C++20 std::atomic<std::shared_ptr<T>>.
+        std::shared_ptr<Context> ctx;
+    };
+
+private:
+    static std::size_t NormalizeSlotCount(std::size_t n)
+    {
+        n = std::max<std::size_t>(n, kMinSlotCount);
+
+        std::size_t power = 1;
+        while (power < n) {
+            if (power > (std::numeric_limits<std::size_t>::max() >> 1)) {
+                return power;
+            }
+            power <<= 1;
         }
-        return Status::OK();
+
+        return power;
+    }
+
+    static std::size_t Hash(TaskId taskId)
+    {
+        return std::hash<TaskId>{}(taskId);
+    }
+
+    // Atomically load shared_ptr<Context> from Slot, ensuring thread safety
+    static std::shared_ptr<Context> AtomicLoadCtx(
+        const Slot& slot,
+        std::memory_order order)
+    {
+        return std::atomic_load_explicit(&slot.ctx, order);
+    }
+
+    // Atomically store task context sharedCtx into slot
+    static void AtomicStoreCtx(
+        Slot& slot,
+        std::shared_ptr<Context> ptr,
+        std::memory_order order)
+    {
+        std::atomic_store_explicit(&slot.ctx, std::move(ptr), order);
     }
 
 private:
     State initialState_;
     std::string taskName_;
     std::atomic<TaskId> nextTaskId_{1};
     // TODO: consider using a lock-free structure !
-    std::mutex mutex_;
+    std::mutex mu_;
     std::unordered_map<TaskId, std::shared_ptr<Context>> tasks_;
+
+    std::vector<Slot> slots_;
+    std::size_t slotMask_{0};
 };
 
-}  // namespace UC::ASU
+}  // namespace UC::ASU