🐛 fix(async_executor): implement cooperative cancellation

lsp-server/src/core/dispatcher.cppm

@@ -1,6 +1,5 @@
 module;
 
-
 export module lsp.core.dispatcher;
 import spdlog;
 
@@ -51,13 +50,13 @@ export namespace lsp::core
     {
     public:
         ExecutionContext(LifecycleCallback lifecycle_callback,
-                         scheduler::AsyncExecutor& scheduler,
+                         scheduler::async_executor::AsyncExecutor& scheduler,
                          manager::ManagerHub& manager_hub) :
             lifecycle_callback_(lifecycle_callback),
             async_executor_(scheduler),
             manager_hub_(manager_hub) {}
 
-        scheduler::AsyncExecutor& GetScheduler() const { return async_executor_; }
+        scheduler::async_executor::AsyncExecutor& GetScheduler() const { return async_executor_; }
         manager::ManagerHub& GetManagerHub() const { return manager_hub_; }
 
         void TriggerLifecycleEvent(ServerLifecycleEvent event) const
@@ -68,7 +67,7 @@ export namespace lsp::core
 
     private:
         LifecycleCallback lifecycle_callback_;
-        scheduler::AsyncExecutor& async_executor_;
+        scheduler::async_executor::AsyncExecutor& async_executor_;
         manager::ManagerHub& manager_hub_;
     };
 
@@ -80,7 +79,7 @@ export namespace lsp::core
         RequestDispatcher();
         ~RequestDispatcher() = default;
 
-        void SetRequestScheduler(scheduler::AsyncExecutor* scheduler);
+        void SetRequestScheduler(scheduler::async_executor::AsyncExecutor* scheduler);
         void SetManagerHub(manager::ManagerHub* manager_hub);
 
         void RegisterRequestProvider(std::shared_ptr<IRequestProvider> provider);
@@ -114,7 +113,7 @@ export namespace lsp::core
 
         LifecycleCallback context_lifecycle_callback_;
 
-        scheduler::AsyncExecutor* async_executor_ = nullptr;
+        scheduler::async_executor::AsyncExecutor* async_executor_ = nullptr;
         manager::ManagerHub* manager_hub_ = nullptr;
     };
 }
@@ -128,7 +127,7 @@ namespace lsp::core
         };
     }
 
-    void RequestDispatcher::SetRequestScheduler(scheduler::AsyncExecutor* scheduler)
+    void RequestDispatcher::SetRequestScheduler(scheduler::async_executor::AsyncExecutor* scheduler)
     {
         async_executor_ = scheduler;
         spdlog::debug("Request scheduler set");

lsp-server/src/core/server.cppm

@@ -75,7 +75,7 @@ export namespace lsp::core
     private:
         RequestDispatcher dispatcher_;
         manager::ManagerHub manager_hub_;
-        scheduler::AsyncExecutor async_executor_;
+        scheduler::async_executor::AsyncExecutor async_executor_;
         std::string interpreter_path_;
 
         std::atomic<bool> is_initialized_ = false;

lsp-server/src/manager/bootstrap.cppm

@@ -12,7 +12,7 @@ export namespace lsp::manager::bootstrap
 {
     void InitializeManagerHub(
         ManagerHub& hub,
-        scheduler::AsyncExecutor& async_executor,
+        scheduler::async_executor::AsyncExecutor& async_executor,
         const std::vector<std::string>& system_lib_paths);
 }
 
@@ -20,7 +20,7 @@ namespace lsp::manager::bootstrap
 {
     void InitializeManagerHub(
         ManagerHub& hub,
-        scheduler::AsyncExecutor& async_executor,
+        scheduler::async_executor::AsyncExecutor& async_executor,
         const std::vector<std::string>& system_lib_paths)
     {
         spdlog::info("Initializing manager hub...");
@@ -31,10 +31,10 @@ namespace lsp::manager::bootstrap
 
             async_executor.Submit(
                 task_name,
-                [&hub, path]() -> std::optional<std::string> {
+                [&hub, path](std::stop_token stop_token) -> std::optional<std::string> {
                     try
                     {
-                        hub.symbols().LoadSystemLibrary(path);
+                        hub.symbols().LoadSystemLibrary(path, stop_token);
                         return std::format("Loaded system library: {}", path);
                     }
                     catch (const std::exception& e)
@@ -43,11 +43,11 @@ namespace lsp::manager::bootstrap
                         throw;
                     }
                 },
-                [path](const std::optional<std::string>& result, bool cancelled) {
-                    if (cancelled)
+                [path](const scheduler::async_executor::TaskResult& result) {
+                    if (result.status == scheduler::async_executor::TaskStatus::kCancelled)
                         spdlog::info("System library load task cancelled: {}", path);
-                    else if (result)
-                        spdlog::info("{}", *result);
+                    else if (result.status == scheduler::async_executor::TaskStatus::kCompleted && result.value)
+                        spdlog::info("{}", *result.value);
                 });
         }
         spdlog::info("Manager hub initialized, system library loading in background");

lsp-server/src/manager/symbol.cppm

@@ -37,11 +37,11 @@ export namespace lsp::manager
         explicit Symbol(EventBus& event_bus);
         ~Symbol();
 
-        void LoadSystemLibrary(const std::string& lib_path);
-        void LoadWorkspace(const protocol::DocumentUri& workspace_uri);
+        void LoadSystemLibrary(const std::string& lib_path, std::stop_token stop_token = {});
+        void LoadWorkspace(const protocol::DocumentUri& workspace_uri, std::stop_token stop_token = {});
 
-        void IndexWorkspaceFiles(const std::vector<protocol::DocumentUri>& uris);
-        void RemoveWorkspaceFiles(const std::vector<protocol::DocumentUri>& uris);
+        void IndexWorkspaceFiles(const std::vector<protocol::DocumentUri>& uris, std::stop_token stop_token = {});
+        void RemoveWorkspaceFiles(const std::vector<protocol::DocumentUri>& uris, std::stop_token stop_token = {});
         void RenameWorkspaceFiles(const std::vector<std::pair<protocol::DocumentUri, protocol::DocumentUri>>& files);
 
         const language::symbol::SymbolTable* GetSymbolTable(const protocol::DocumentUri& uri) const;
@@ -306,11 +306,16 @@ namespace lsp::manager
 
     Symbol::~Symbol() = default;
 
-    void Symbol::LoadSystemLibrary(const std::string& lib_path)
+    void Symbol::LoadSystemLibrary(const std::string& lib_path, std::stop_token stop_token)
     {
         spdlog::info("Loading system library from: {}", lib_path);
         auto start = std::chrono::steady_clock::now();
 
+        if (stop_token.stop_requested())
+        {
+            return;
+        }
+
         if (!std::filesystem::exists(lib_path))
         {
             spdlog::warn("System library path does not exist: {}", lib_path);
@@ -326,6 +331,11 @@ namespace lsp::manager
 
         for (const auto& entry : std::filesystem::recursive_directory_iterator(lib_path, options))
         {
+            if (stop_token.stop_requested())
+            {
+                return;
+            }
+
             if (!entry.is_regular_file())
                 continue;
 
@@ -338,6 +348,10 @@ namespace lsp::manager
             spdlog::trace("Indexing library file: {}", entry.path().string());
 
             auto table = BuildSymbolTableFromFile(entry.path());
+            if (stop_token.stop_requested())
+            {
+                return;
+            }
             if (!table)
             {
                 spdlog::trace("Failed to build symbol table for: {}", entry.path().string());
@@ -362,6 +376,11 @@ namespace lsp::manager
             ++loaded;
         }
 
+        if (stop_token.stop_requested())
+        {
+            return;
+        }
+
         {
             std::unique_lock<std::shared_mutex> lock(mutex_);
             system_symbols_ = std::move(new_symbols);
@@ -378,12 +397,17 @@ namespace lsp::manager
                      duration);
     }
 
-    void Symbol::LoadWorkspace(const protocol::DocumentUri& workspace_uri)
+    void Symbol::LoadWorkspace(const protocol::DocumentUri& workspace_uri, std::stop_token stop_token)
     {
         auto workspace_path = UriToPath(workspace_uri);
         spdlog::info("Loading workspace from: {}", workspace_path);
         auto start = std::chrono::steady_clock::now();
 
+        if (stop_token.stop_requested())
+        {
+            return;
+        }
+
         if (!std::filesystem::exists(workspace_path))
         {
             spdlog::warn("Workspace path does not exist: {}", workspace_path);
@@ -399,6 +423,11 @@ namespace lsp::manager
 
         for (const auto& entry : std::filesystem::recursive_directory_iterator(workspace_path, options))
         {
+            if (stop_token.stop_requested())
+            {
+                return;
+            }
+
             if (!entry.is_regular_file())
                 continue;
 
@@ -407,6 +436,10 @@ namespace lsp::manager
                 continue;
 
             auto table = BuildSymbolTableFromFile(entry.path());
+            if (stop_token.stop_requested())
+            {
+                return;
+            }
             if (!table)
             {
                 ++failed;
@@ -431,6 +464,11 @@ namespace lsp::manager
             ++loaded;
         }
 
+        if (stop_token.stop_requested())
+        {
+            return;
+        }
+
         {
             std::unique_lock<std::shared_mutex> lock(mutex_);
             workspace_symbols_ = std::move(new_symbols);
@@ -447,7 +485,7 @@ namespace lsp::manager
                      duration);
     }
 
-    void Symbol::IndexWorkspaceFiles(const std::vector<protocol::DocumentUri>& uris)
+    void Symbol::IndexWorkspaceFiles(const std::vector<protocol::DocumentUri>& uris, std::stop_token stop_token)
     {
         std::unordered_map<std::string, StoredSymbolEntry> updates;
         std::vector<std::string> removals;
@@ -457,6 +495,11 @@ namespace lsp::manager
 
         for (const auto& uri : uris)
         {
+            if (stop_token.stop_requested())
+            {
+                return;
+            }
+
             auto file_path = std::filesystem::path(UriToPath(uri));
             auto kind = GetTslFileKind(file_path);
             if (kind == TslFileKind::kOther)
@@ -473,6 +516,10 @@ namespace lsp::manager
             }
 
             auto table = BuildSymbolTableFromFile(file_path);
+            if (stop_token.stop_requested())
+            {
+                return;
+            }
             if (!table)
             {
                 removals.push_back(std::move(normalized_uri));
@@ -494,7 +541,7 @@ namespace lsp::manager
             updates[normalized_uri] = std::move(stored);
         }
 
-        if (updates.empty() && removals.empty())
+        if (stop_token.stop_requested() || (updates.empty() && removals.empty()))
         {
             return;
         }
@@ -513,7 +560,7 @@ namespace lsp::manager
         }
     }
 
-    void Symbol::RemoveWorkspaceFiles(const std::vector<protocol::DocumentUri>& uris)
+    void Symbol::RemoveWorkspaceFiles(const std::vector<protocol::DocumentUri>& uris, std::stop_token stop_token)
     {
         if (uris.empty())
         {
@@ -525,6 +572,11 @@ namespace lsp::manager
 
         for (const auto& uri : uris)
         {
+            if (stop_token.stop_requested())
+            {
+                return;
+            }
+
             auto file_path = std::filesystem::path(UriToPath(uri));
             auto kind = GetTslFileKind(file_path);
             if (kind == TslFileKind::kOther)
@@ -534,7 +586,7 @@ namespace lsp::manager
             removals.push_back(PathToUri(file_path));
         }
 
-        if (removals.empty())
+        if (stop_token.stop_requested() || removals.empty())
         {
             return;
         }

lsp-server/src/provider/initialize/initialize.cppm

@@ -39,9 +39,6 @@ export namespace lsp::provider
 
 namespace lsp::provider
 {
-    
-
-    
 
     std::string Initialize::ProvideResponse(const protocol::RequestMessage& request, ExecutionContext& context)
     {
@@ -219,24 +216,24 @@ namespace lsp::provider
         for (const auto& workspace_folder : workspace_folders)
         {
             auto task_id = std::format("Load workspace symbols: {}", workspace_folder.uri);
-            scheduler.Submit(task_id, [&manager_hub, uri = workspace_folder.uri, folder_name = workspace_folder.name]() -> std::optional<std::string> {
+            scheduler.Submit(task_id, [&manager_hub, uri = workspace_folder.uri, folder_name = workspace_folder.name](std::stop_token stop_token) -> std::optional<std::string> {
                 try
                 {
-                    manager_hub.symbols().LoadWorkspace(uri);
+                    manager_hub.symbols().LoadWorkspace(uri, stop_token);
                     return std::format("Loaded workspace {} symbols", uri);
                 }
                 catch (const std::exception& e)
                 {
                     spdlog::error("Failed to load workspace {} symbols: {}", folder_name, e.what());
                     throw; // Request会处理异常
-                } }, [](const std::optional<std::string>& result, bool cancelled) {
-                if (cancelled)
+                } }, [](const ::lsp::scheduler::async_executor::TaskResult& result) {
+                if (result.status == ::lsp::scheduler::async_executor::TaskStatus::kCancelled)
                 {
                     spdlog::debug("Workspace loading task cancelled");
                 }
-                else if (result)
+                else if (result.status == ::lsp::scheduler::async_executor::TaskStatus::kCompleted && result.value)
                 {
-                    spdlog::info("Workspace loading result: {}", *result);
+                    spdlog::info("Workspace loading result: {}", *result.value);
                 } });
         }
 

lsp-server/src/provider/workspace/did_change_workspace_folders.cppm

@@ -1,6 +1,5 @@
 module;
 
-
 export module lsp.provider.workspace.did_change_workspace_folders;
 import spdlog;
 
@@ -9,6 +8,7 @@ import std;
 import lsp.protocol;
 import lsp.codec.facade;
 import lsp.provider.base.interface;
+import lsp.scheduler.async_executor;
 
 export namespace lsp::provider::workspace
 {
@@ -123,10 +123,17 @@ namespace lsp::provider::workspace
             return folders;
         }
 
-        std::vector<protocol::DocumentUri> EnumerateWorkspaceFiles(const protocol::DocumentUri& workspace_uri)
+        std::vector<protocol::DocumentUri> EnumerateWorkspaceFiles(
+            const protocol::DocumentUri& workspace_uri,
+            std::stop_token stop_token)
         {
             std::vector<protocol::DocumentUri> uris;
 
+            if (stop_token.stop_requested())
+            {
+                return uris;
+            }
+
             std::filesystem::path workspace_path;
             try
             {
@@ -147,6 +154,11 @@ namespace lsp::provider::workspace
 
             for (const auto& entry : std::filesystem::recursive_directory_iterator(workspace_path, options))
             {
+                if (stop_token.stop_requested())
+                {
+                    return {};
+                }
+
                 if (!entry.is_regular_file())
                 {
                     continue;
@@ -211,22 +223,26 @@ namespace lsp::provider::workspace
             auto task_id = std::format("Remove workspace folder: {}", folder.uri);
             scheduler.Submit(
                 task_id,
-                [&hub, uri = folder.uri]() -> std::optional<std::string> {
-                    auto uris = EnumerateWorkspaceFiles(uri);
+                [&hub, uri = folder.uri](std::stop_token stop_token) -> std::optional<std::string> {
+                    auto uris = EnumerateWorkspaceFiles(uri, stop_token);
+                    if (stop_token.stop_requested())
+                    {
+                        return std::nullopt;
+                    }
                     if (!uris.empty())
                     {
-                        hub.symbols().RemoveWorkspaceFiles(uris);
+                        hub.symbols().RemoveWorkspaceFiles(uris, stop_token);
                     }
                     return std::format("Removed {} workspace file(s)", uris.size());
                 },
-                [](const std::optional<std::string>& result, bool cancelled) {
-                    if (cancelled)
+                [](const ::lsp::scheduler::async_executor::TaskResult& result) {
+                    if (result.status == ::lsp::scheduler::async_executor::TaskStatus::kCancelled)
                     {
                         spdlog::debug("Workspace folder removal task cancelled");
                     }
-                    else if (result)
+                    else if (result.status == ::lsp::scheduler::async_executor::TaskStatus::kCompleted && result.value)
                     {
-                        spdlog::info("{}", *result);
+                        spdlog::info("{}", *result.value);
                     }
                 });
         }
@@ -236,22 +252,26 @@ namespace lsp::provider::workspace
             auto task_id = std::format("Index workspace folder: {}", folder.uri);
             scheduler.Submit(
                 task_id,
-                [&hub, uri = folder.uri]() -> std::optional<std::string> {
-                    auto uris = EnumerateWorkspaceFiles(uri);
+                [&hub, uri = folder.uri](std::stop_token stop_token) -> std::optional<std::string> {
+                    auto uris = EnumerateWorkspaceFiles(uri, stop_token);
+                    if (stop_token.stop_requested())
+                    {
+                        return std::nullopt;
+                    }
                     if (!uris.empty())
                     {
-                        hub.symbols().IndexWorkspaceFiles(uris);
+                        hub.symbols().IndexWorkspaceFiles(uris, stop_token);
                     }
                     return std::format("Indexed {} workspace file(s)", uris.size());
                 },
-                [](const std::optional<std::string>& result, bool cancelled) {
-                    if (cancelled)
+                [](const ::lsp::scheduler::async_executor::TaskResult& result) {
+                    if (result.status == ::lsp::scheduler::async_executor::TaskStatus::kCancelled)
                     {
                         spdlog::debug("Workspace folder indexing task cancelled");
                     }
-                    else if (result)
+                    else if (result.status == ::lsp::scheduler::async_executor::TaskStatus::kCompleted && result.value)
                     {
-                        spdlog::info("{}", *result);
+                        spdlog::info("{}", *result.value);
                     }
                 });
         }

lsp-server/src/provider/workspace/execute_command.cppm

@@ -1,6 +1,5 @@
 module;
 
-
 export module lsp.provider.workspace.execute_command;
 import spdlog;
 
@@ -107,8 +106,8 @@ namespace lsp::provider::workspace
             auto& hub = context.GetManagerHub();
             auto& scheduler = context.GetScheduler();
             auto task_id = std::format("ExecuteCommand load workspace: {}", *uri);
-            scheduler.Submit(task_id, [&hub, uri = *uri]() -> std::optional<std::string> {
-                hub.symbols().LoadWorkspace(uri);
+            scheduler.Submit(task_id, [&hub, uri = *uri](std::stop_token stop_token) -> std::optional<std::string> {
+                hub.symbols().LoadWorkspace(uri, stop_token);
                 return std::string("ok");
             });
 
@@ -125,8 +124,8 @@ namespace lsp::provider::workspace
             auto& hub = context.GetManagerHub();
             auto& scheduler = context.GetScheduler();
             auto task_id = std::format("ExecuteCommand index files: {}", uris.size());
-            scheduler.Submit(task_id, [&hub, uris = std::move(uris)]() -> std::optional<std::string> {
-                hub.symbols().IndexWorkspaceFiles(uris);
+            scheduler.Submit(task_id, [&hub, uris = std::move(uris)](std::stop_token stop_token) -> std::optional<std::string> {
+                hub.symbols().IndexWorkspaceFiles(uris, stop_token);
                 return std::string("ok");
             });
 

lsp-server/src/scheduler/async_executor.cppm

@@ -4,70 +4,61 @@ import spdlog;
 import taskflow;
 import std;
 
-export namespace lsp::scheduler
+export namespace lsp::scheduler::async_executor
 {
-    class AsyncExecutor;
-
-    namespace detail
+    enum class TaskStatus
     {
-        struct ExecutionState
-        {
-            std::atomic<bool> cancelled{ false };
-            std::mutex mutex;
-            std::condition_variable cv;
-            bool completed = false;
-            bool callback_completed = false;
-            std::optional<std::string> result;
-            std::exception_ptr error;
-            std::chrono::steady_clock::time_point start_time{};
-        };
+        kCompleted,
+        kCancelled,
+        kFailed,
+    };
 
-        struct ActiveEntry
-        {
-            std::shared_ptr<ExecutionState> state;
-            std::function<void(const std::optional<std::string>& result, bool cancelled)> callback;
-            std::chrono::steady_clock::time_point start_time;
-        };
+    struct TaskResult
+    {
+        TaskStatus status;
+        std::optional<std::string> value;
+        std::exception_ptr error;
+    };
 
-        struct ExecutorMetrics
-        {
-            std::size_t running = 0;
-            std::uint64_t submitted = 0;
-            std::uint64_t completed = 0;
-            std::uint64_t failed = 0;
-            std::uint64_t cancelled = 0;
-        };
-    }
+    struct ExecutorMetrics
+    {
+        std::size_t running = 0;
+        std::uint64_t submitted = 0;
+        std::uint64_t completed = 0;
+        std::uint64_t failed = 0;
+        std::uint64_t cancelled = 0;
+    };
+
+    class AsyncExecutor;
 
     class TaskHandle
     {
     public:
         TaskHandle() = default;
-        TaskHandle(AsyncExecutor* executor, std::string task_id, std::shared_ptr<detail::ExecutionState> state) :
-            executor_(executor), task_id_(std::move(task_id)), state_(std::move(state))
-        {
-        }
-
-        bool Valid() const
-        {
-            return executor_ != nullptr && !task_id_.empty() && !state_.expired();
-        }
 
+        bool Valid() const;
         bool Cancel() const;
-        bool Wait() const;
-        std::optional<std::string> GetResult() const;
+        std::optional<TaskResult> Wait() const;
+        std::optional<TaskResult> TryGetResult() const;
 
     private:
-        AsyncExecutor* executor_ = nullptr;
-        std::string task_id_;
-        std::weak_ptr<detail::ExecutionState> state_;
+        struct State;
+
+        explicit TaskHandle(std::shared_ptr<State> state);
+
+        std::shared_ptr<State> state_;
+
+        friend class AsyncExecutor;
     };
 
     class AsyncExecutor
     {
     public:
-        using TaskClosure = std::function<std::optional<std::string>()>;
-        using TaskCallback = std::function<void(const std::optional<std::string>&, bool)>;
+        using TaskClosure = std::function<std::optional<std::string>(std::stop_token)>;
+
+        // A callback must not wait on its own handle or call WaitAll() on this
+        // executor because either operation would wait for the callback itself.
+        using TaskCallback = std::function<void(const TaskResult&)>;
 
         explicit AsyncExecutor(std::size_t concurrency = std::thread::hardware_concurrency());
         ~AsyncExecutor();
@@ -77,89 +68,114 @@ export namespace lsp::scheduler
 
         TaskHandle Submit(const std::string& task_id, TaskClosure task, TaskCallback callback = nullptr);
         bool Cancel(const std::string& task_id);
-        bool WaitForTask(const std::string& task_id);
         void WaitAll();
 
         std::size_t GetRunningTaskCount() const;
         void LogStatus() const;
-        detail::ExecutorMetrics GetStatistics() const;
+        ExecutorMetrics GetStatistics() const;
 
     private:
-        void ExecuteTask(const std::string& task_id, TaskClosure task, std::shared_ptr<detail::ExecutionState> state, TaskCallback callback);
-        void CompleteTask(const std::string& task_id,
-                          std::chrono::steady_clock::time_point start_time,
-                          std::shared_ptr<detail::ExecutionState> state,
-                          const std::optional<std::string>& result,
-                          bool cancelled,
-                          bool failed,
-                          TaskCallback callback);
-        bool RegisterTask(const std::string& task_id, const detail::ActiveEntry& ctx);
-        void UnregisterTask(const std::string& task_id, const std::shared_ptr<detail::ExecutionState>& state);
-        bool MarkTaskCancelled(const std::string& task_id);
-        std::chrono::milliseconds GetElapsedTime(std::chrono::steady_clock::time_point start_time) const;
+        using TaskState = TaskHandle::State;
+
+        static bool RequestCancel(const std::shared_ptr<TaskState>& state);
+        void ExecuteTask(std::shared_ptr<TaskState> state, TaskClosure task, TaskCallback callback);
+        void CompleteTask(const std::shared_ptr<TaskState>& state, TaskResult result, const TaskCallback& callback);
+        void FinalizeTask(const std::shared_ptr<TaskState>& state);
 
-    private:
         tf::Executor executor_;
         mutable std::mutex mutex_;
-        std::unordered_map<std::string, detail::ActiveEntry> running_tasks_;
+        std::condition_variable active_tasks_cv_;
+        std::unordered_map<std::string, std::shared_ptr<TaskState>> current_tasks_;
+        std::unordered_set<std::shared_ptr<TaskState>> active_tasks_;
 
         std::atomic<std::uint64_t> submitted_{ 0 };
         std::atomic<std::uint64_t> completed_{ 0 };
         std::atomic<std::uint64_t> failed_{ 0 };
         std::atomic<std::uint64_t> cancelled_{ 0 };
 
-        static constexpr std::size_t kStatusLogInterval = 10;
+        friend class TaskHandle;
     };
 }
 
-namespace lsp::scheduler
+namespace lsp::scheduler::async_executor
 {
+    struct TaskHandle::State
+    {
+        enum class Phase
+        {
+            kPending,
+            kRunning,
+            kCompleted,
+        };
+
+        explicit State(std::string id) : task_id(std::move(id)) {}
+
+        std::string task_id;
+        std::chrono::steady_clock::time_point start_time = std::chrono::steady_clock::now();
+        std::stop_source stop_source;
+        mutable std::mutex mutex;
+        std::condition_variable cv;
+        Phase phase = Phase::kPending;
+        bool cancel_requested = false;
+        std::optional<TaskResult> result;
+        bool callback_completed = false;
+    };
+
     namespace
     {
         constexpr const char* kLogTag = "AsyncExecutor";
 
+        std::size_t NormalizeConcurrency(std::size_t concurrency)
+        {
+            return std::max<std::size_t>(1, concurrency);
+        }
+
         template<typename Duration>
         std::string FormatDuration(Duration duration)
         {
-            auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();
-            if (ms < 1000)
-                return std::to_string(ms) + "ms";
-            else
-                return std::to_string(ms / 1000.0) + "s";
+            auto milliseconds = std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();
+            if (milliseconds < 1000)
+                return std::to_string(milliseconds) + "ms";
+            return std::to_string(milliseconds / 1000.0) + "s";
         }
     }
 
+    TaskHandle::TaskHandle(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+    bool TaskHandle::Valid() const
+    {
+        return state_ != nullptr;
+    }
+
     bool TaskHandle::Cancel() const
     {
-        if (!Valid())
-            return false;
-        return executor_->Cancel(task_id_);
+        return AsyncExecutor::RequestCancel(state_);
     }
 
-    bool TaskHandle::Wait() const
+    std::optional<TaskResult> TaskHandle::Wait() const
     {
-        auto state = state_.lock();
-        if (!state)
-            return false;
-        std::unique_lock<std::mutex> lk(state->mutex);
-        state->cv.wait(lk, [state]() { return state->completed && state->callback_completed; });
-        return true;
-    }
-
-    std::optional<std::string> TaskHandle::GetResult() const
-    {
-        auto state = state_.lock();
-        if (!state)
+        if (!state_)
             return std::nullopt;
-        std::unique_lock<std::mutex> lk(state->mutex);
-        if (!state->completed || !state->callback_completed || state->error)
-            return std::nullopt;
-        return state->result;
+
+        std::unique_lock lock(state_->mutex);
+        state_->cv.wait(lock, [this] { return state_->callback_completed; });
+        return state_->result;
     }
 
-    AsyncExecutor::AsyncExecutor(std::size_t concurrency) : executor_(concurrency)
+    std::optional<TaskResult> TaskHandle::TryGetResult() const
     {
-        spdlog::info("{}: Initialized with {} threads", kLogTag, concurrency);
+        if (!state_)
+            return std::nullopt;
+
+        std::lock_guard lock(state_->mutex);
+        if (!state_->callback_completed)
+            return std::nullopt;
+        return state_->result;
+    }
+
+    AsyncExecutor::AsyncExecutor(std::size_t concurrency) : executor_(NormalizeConcurrency(concurrency))
+    {
+        spdlog::info("{}: Initialized with {} threads", kLogTag, NormalizeConcurrency(concurrency));
     }
 
     AsyncExecutor::~AsyncExecutor()
@@ -171,91 +187,71 @@ namespace lsp::scheduler
 
     TaskHandle AsyncExecutor::Submit(const std::string& task_id, TaskClosure task, TaskCallback callback)
     {
-        detail::ActiveEntry entry;
-        entry.state = std::make_shared<detail::ExecutionState>();
-        entry.state->start_time = std::chrono::steady_clock::now();
-        entry.callback = callback;
-        entry.start_time = entry.state->start_time;
+        auto state = std::make_shared<TaskState>(task_id);
+        std::shared_ptr<TaskState> replaced_state;
+        {
+            std::lock_guard lock(mutex_);
+            auto current = current_tasks_.find(task_id);
+            if (current != current_tasks_.end())
+                replaced_state = current->second;
+            current_tasks_[task_id] = state;
+            active_tasks_.insert(state);
+            ++submitted_;
+        }
 
-        bool replaced_existing = RegisterTask(task_id, entry);
-
-        if (replaced_existing)
+        if (replaced_state)
+        {
+            RequestCancel(replaced_state);
             spdlog::warn("[{}] Replaced existing task", task_id);
+        }
 
         try
         {
-            executor_.async([this, task_id, task = std::move(task), state = entry.state, callback]() mutable {
-                ExecuteTask(task_id, std::move(task), state, callback);
+            executor_.async([this, state, task = std::move(task), callback]() mutable {
+                ExecuteTask(std::move(state), std::move(task), std::move(callback));
             });
         }
-        catch (const std::exception& e)
+        catch (...)
         {
-            spdlog::error("[{}] Failed to submit task: {}", task_id, e.what());
-            CompleteTask(task_id, entry.start_time, entry.state, std::nullopt, false, true, callback);
+            auto error = std::current_exception();
+            spdlog::error("[{}] Failed to submit task", task_id);
+            CompleteTask(state, TaskResult{ TaskStatus::kFailed, std::nullopt, error }, callback);
         }
 
-        return TaskHandle(this, task_id, entry.state);
+        return TaskHandle(std::move(state));
     }
 
     bool AsyncExecutor::Cancel(const std::string& task_id)
     {
-        return MarkTaskCancelled(task_id);
-    }
-
-    bool AsyncExecutor::WaitForTask(const std::string& task_id)
-    {
-        std::shared_ptr<detail::ExecutionState> state;
+        std::shared_ptr<TaskState> state;
         {
-            std::unique_lock<std::mutex> lock(mutex_);
-            auto it = running_tasks_.find(task_id);
-            if (it == running_tasks_.end())
+            std::lock_guard lock(mutex_);
+            auto current = current_tasks_.find(task_id);
+            if (current == current_tasks_.end())
                 return false;
-            state = it->second.state;
+            state = current->second;
         }
-
-        std::unique_lock<std::mutex> lk(state->mutex);
-        state->cv.wait(lk, [state]() { return state->completed && state->callback_completed; });
-        return true;
+        return RequestCancel(state);
     }
 
     void AsyncExecutor::WaitAll()
     {
-        std::vector<std::shared_ptr<detail::ExecutionState>> tasks;
         {
-            std::unique_lock<std::mutex> lock(mutex_);
-            for (const auto& [_, entry] : running_tasks_)
-            {
-                tasks.push_back(entry.state);
-            }
+            std::unique_lock lock(mutex_);
+            active_tasks_cv_.wait(lock, [this] { return active_tasks_.empty(); });
         }
-
-        for (const auto& state : tasks)
-        {
-            std::unique_lock<std::mutex> lk(state->mutex);
-            state->cv.wait(lk, [state]() { return state->completed && state->callback_completed; });
-        }
-
-        // Ensure the underlying executor finishes any tasks that may have been
-        // cancelled or replaced but are still running.
         executor_.wait_for_all();
     }
 
     std::size_t AsyncExecutor::GetRunningTaskCount() const
     {
-        std::unique_lock<std::mutex> lock(mutex_);
-        return running_tasks_.size();
+        std::lock_guard lock(mutex_);
+        return active_tasks_.size();
     }
 
     void AsyncExecutor::LogStatus() const
     {
-        std::unique_lock<std::mutex> lock(mutex_);
-        detail::ExecutorMetrics metrics;
-        metrics.running = running_tasks_.size();
-        metrics.submitted = submitted_;
-        metrics.completed = completed_;
-        metrics.failed = failed_;
-        metrics.cancelled = cancelled_;
-
+        const auto metrics = GetStatistics();
         spdlog::info("[{}] Running: {}, Submitted: {}, Completed: {}, Failed: {}, Cancelled: {}",
                      kLogTag,
                      metrics.running,
@@ -265,157 +261,131 @@ namespace lsp::scheduler
                      metrics.cancelled);
     }
 
-    detail::ExecutorMetrics AsyncExecutor::GetStatistics() const
+    ExecutorMetrics AsyncExecutor::GetStatistics() const
     {
-        detail::ExecutorMetrics metrics;
-        std::unique_lock<std::mutex> lock(mutex_);
-        metrics.running = running_tasks_.size();
-        metrics.submitted = submitted_;
-        metrics.completed = completed_;
-        metrics.failed = failed_;
-        metrics.cancelled = cancelled_;
+        ExecutorMetrics metrics;
+        {
+            std::lock_guard lock(mutex_);
+            metrics.running = active_tasks_.size();
+        }
+        metrics.submitted = submitted_.load();
+        metrics.completed = completed_.load();
+        metrics.failed = failed_.load();
+        metrics.cancelled = cancelled_.load();
         return metrics;
     }
 
-    void AsyncExecutor::ExecuteTask(const std::string& task_id, TaskClosure task, std::shared_ptr<detail::ExecutionState> state, TaskCallback callback)
+    bool AsyncExecutor::RequestCancel(const std::shared_ptr<TaskState>& state)
     {
+        if (!state)
+            return false;
+
+        {
+            std::lock_guard lock(state->mutex);
+            if (state->phase == TaskState::Phase::kCompleted || state->cancel_requested)
+                return false;
+            state->cancel_requested = true;
+        }
+        return state->stop_source.request_stop();
+    }
+
+    void AsyncExecutor::ExecuteTask(std::shared_ptr<TaskState> state, TaskClosure task, TaskCallback callback)
+    {
+        const auto stop_token = state->stop_source.get_token();
+        bool cancelled_before_start = false;
+        {
+            std::lock_guard lock(state->mutex);
+            cancelled_before_start = state->cancel_requested;
+            if (!cancelled_before_start)
+                state->phase = TaskState::Phase::kRunning;
+        }
+
+        if (cancelled_before_start)
+        {
+            CompleteTask(state, TaskResult{ TaskStatus::kCancelled, std::nullopt, nullptr }, callback);
+            return;
+        }
+
+        TaskResult result;
         try
         {
-            auto result = task();
-            CompleteTask(task_id,
-                         state->start_time,
-                         state,
-                         result,
-                         state->cancelled.load(std::memory_order_relaxed),
-                         false,
-                         callback);
+            auto value = task(stop_token);
+            if (stop_token.stop_requested())
+                result = TaskResult{ TaskStatus::kCancelled, std::nullopt, nullptr };
+            else
+                result = TaskResult{ TaskStatus::kCompleted, std::move(value), nullptr };
         }
         catch (const std::exception& e)
         {
-            spdlog::error("[{}] Task threw exception: {}", task_id, e.what());
-            CompleteTask(task_id,
-                         state->start_time,
-                         state,
-                         std::nullopt,
-                         state->cancelled.load(std::memory_order_relaxed),
-                         true,
-                         callback);
+            spdlog::error("[{}] Task threw exception: {}", state->task_id, e.what());
+            result = TaskResult{ TaskStatus::kFailed, std::nullopt, std::current_exception() };
         }
         catch (...)
         {
-            spdlog::error("[{}] Task threw unknown exception", task_id);
-            CompleteTask(task_id,
-                         state->start_time,
-                         state,
-                         std::nullopt,
-                         state->cancelled.load(std::memory_order_relaxed),
-                         true,
-                         callback);
+            spdlog::error("[{}] Task threw unknown exception", state->task_id);
+            result = TaskResult{ TaskStatus::kFailed, std::nullopt, std::current_exception() };
         }
+
+        CompleteTask(state, std::move(result), callback);
     }
 
-    void AsyncExecutor::CompleteTask(const std::string& task_id,
-                                     std::chrono::steady_clock::time_point start_time,
-                                     std::shared_ptr<detail::ExecutionState> state,
-                                     const std::optional<std::string>& result,
-                                     bool cancelled,
-                                     bool failed,
-                                     TaskCallback callback)
+    void AsyncExecutor::CompleteTask(const std::shared_ptr<TaskState>& state, TaskResult result, const TaskCallback& callback)
     {
-        if (!state)
-            return;
-
-        const bool is_cancelled = cancelled || state->cancelled.load(std::memory_order_relaxed);
         {
-            std::unique_lock<std::mutex> lk(state->mutex);
-            state->completed = true;
+            std::lock_guard lock(state->mutex);
+            if (result.status == TaskStatus::kCompleted && state->cancel_requested)
+                result = TaskResult{ TaskStatus::kCancelled, std::nullopt, nullptr };
+            state->phase = TaskState::Phase::kCompleted;
             state->result = result;
-            state->cancelled = is_cancelled;
-            if (failed)
-                state->error = std::make_exception_ptr(std::runtime_error("Task failed"));
         }
 
-        UnregisterTask(task_id, state);
-
-        auto elapsed = GetElapsedTime(start_time);
-        bool callback_failed = false;
         if (callback)
         {
             try
             {
-                callback(result, is_cancelled);
+                callback(result);
             }
             catch (...)
             {
-                callback_failed = true;
-                std::unique_lock<std::mutex> lk(state->mutex);
-                state->error = std::current_exception();
+                spdlog::error("[{}] Task callback threw exception", state->task_id);
             }
         }
 
-        const bool has_failed = failed || callback_failed;
-        if (has_failed)
-            ++failed_;
-        else if (is_cancelled)
-            ++cancelled_;
-        else
-            ++completed_;
-
+        switch (result.status)
         {
-            std::unique_lock<std::mutex> lk(state->mutex);
+        case TaskStatus::kCompleted:
+            ++completed_;
+            break;
+        case TaskStatus::kCancelled:
+            ++cancelled_;
+            break;
+        case TaskStatus::kFailed:
+            ++failed_;
+            break;
+        }
+
+        FinalizeTask(state);
+
+        const auto elapsed = std::chrono::steady_clock::now() - state->start_time;
+        spdlog::info("[{}] Task completed. status={}, elapsed={}",
+                     state->task_id,
+                     static_cast<int>(result.status),
+                     FormatDuration(elapsed));
+    }
+
+    void AsyncExecutor::FinalizeTask(const std::shared_ptr<TaskState>& state)
+    {
+        {
+            std::lock_guard lock(mutex_);
+            active_tasks_.erase(state);
+            auto current = current_tasks_.find(state->task_id);
+            if (current != current_tasks_.end() && current->second == state)
+                current_tasks_.erase(current);
+
+            std::lock_guard state_lock(state->mutex);
             state->callback_completed = true;
         }
         state->cv.notify_all();
-
-        if (callback_failed)
-            spdlog::error("[{}] Task callback threw exception", task_id);
-
-        spdlog::info("[{}] Task completed. cancelled={}, failed={}, elapsed={}", task_id, is_cancelled, has_failed, FormatDuration(elapsed));
+        active_tasks_cv_.notify_all();
     }
-
-    bool AsyncExecutor::RegisterTask(const std::string& task_id, const detail::ActiveEntry& ctx)
-    {
-        std::unique_lock<std::mutex> lock(mutex_);
-        auto it = running_tasks_.find(task_id);
-        const bool replaced = it != running_tasks_.end();
-        if (it != running_tasks_.end())
-        {
-            // Mark prior task as cancelled; completion will observe this flag.
-            it->second.state->cancelled.store(true, std::memory_order_relaxed);
-        }
-        running_tasks_[task_id] = ctx;
-        ++submitted_;
-        return replaced;
-    }
-
-    void AsyncExecutor::UnregisterTask(const std::string& task_id, const std::shared_ptr<detail::ExecutionState>& state)
-    {
-        std::unique_lock<std::mutex> lock(mutex_);
-        auto it = running_tasks_.find(task_id);
-        if (it != running_tasks_.end() && it->second.state == state)
-        {
-            running_tasks_.erase(it);
-        }
-    }
-
-    bool AsyncExecutor::MarkTaskCancelled(const std::string& task_id)
-    {
-        std::unique_lock<std::mutex> lock(mutex_);
-        auto it = running_tasks_.find(task_id);
-        if (it == running_tasks_.end())
-            return false;
-
-        auto& state = it->second.state;
-        {
-            std::unique_lock<std::mutex> lk(state->mutex);
-            state->cancelled = true;
-        }
-        return true;
-    }
-
-    std::chrono::milliseconds AsyncExecutor::GetElapsedTime(std::chrono::steady_clock::time_point start_time) const
-    {
-        return std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - start_time);
-    }
-
 }