refactor(server): interest-counted module compile cancellation

src/server/compiler/compile_graph.h

@@ -3,15 +3,35 @@
 #include <cstdint>
 #include <functional>
 #include <memory>
+#include <optional>
 
 #include "kota/async/async.h"
 #include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace clice {
 
 struct CompileUnit {
+    /// Result of one compilation round, observed by waiters after the
+    /// round's completion event fires.
+    enum class Outcome : std::uint8_t {
+        /// The round was cancelled (file update or loss of interest) and its
+        /// result discarded; waiters that still hold interest retry.
+        Stale,
+        Success,
+        /// Dispatch failed or a dependency cycle was detected; waiters
+        /// propagate the failure instead of retrying.
+        Failed,
+    };
+
+    /// State of one compilation round. Waiters capture the shared_ptr before
+    /// suspending so the completion event outlives both map mutations and
+    /// the unit task that publishes the outcome.
+    struct Round {
+        kota::event completion;
+        Outcome outcome = Outcome::Stale;
+    };
+
     std::uint32_t path_id = 0;
 
     /// Dependencies discovered lazily by resolve_fn.
@@ -26,15 +46,52 @@ struct CompileUnit {
     bool dirty = true;
     bool compiling = false;
 
-    /// Monotonic counter bumped by update(); used by compile_impl to detect
-    /// stale completions without ABA risk from raw-pointer comparison.
+    /// In-flight interest count: one per requesting root plus one per edge
+    /// held by a dependent unit's running task. Staying at zero while
+    /// compiling cancels this unit's round. Not a lifetime count.
+    std::uint32_t refcount = 0;
+
+    /// A zero-interest cancellation check is already queued for this unit.
+    bool zero_check_pending = false;
+
+    /// Monotonic counter bumped by update(); detects results that became
+    /// stale while the cancellation raced with a dispatch completion.
     std::uint64_t generation = 0;
 
+    /// Cancellation scope of the current round; replaced after every cancel
+    /// so the next round always starts with a fresh token.
     std::unique_ptr<kota::cancellation_source> source =
         std::make_unique<kota::cancellation_source>();
-    std::unique_ptr<kota::event> completion;
+
+    /// Current (or most recent) compilation round.
+    std::shared_ptr<Round> round;
 };
 
+/// Module compilation DAG with interest-counted cancellation.
+///
+/// Each dirty unit is compiled by an independent task spawned into the
+/// graph's task group, cancellable only through its own round token —
+/// requesters and dependent units merely wait on the round's completion
+/// event. Lifecycle by phase:
+///
+/// - Request arrival: compile()/compile_deps() takes a root reference on the
+///   requested unit(s) and waits. A dirty unit with no running round gets a
+///   unit task spawned; the task acquires edge references on its direct
+///   dependencies, waits for them, then dispatches its own compilation.
+/// - Request cancel: the requester's frame unwinds and drops its root
+///   reference. A unit whose interest stays at zero for an event-loop tick
+///   has its round cancelled; the exiting task drops its edge references,
+///   cascading the cancellation through no-longer-needed dependencies.
+///   Shared dependencies keep compiling as long as any other consumer holds
+///   interest, and a transient drop (a retry re-acquiring a retained
+///   dependency after re-resolve) does not disturb the running compilation.
+/// - File update: update() marks the unit and its transitive dependents
+///   dirty and cancels their in-flight rounds — the results are stale.
+///   Interest is NOT touched; waiters observe the stale round and drive a
+///   fresh one with the new content.
+/// - Compile finish: the unit task publishes success/failure through its
+///   round and wakes waiters; success clears dirty, failure (compile error,
+///   dependency cycle) propagates to waiters without retry.
 class CompileGraph {
 public:
     /// Performs the actual compilation (e.g. produce PCM file).
@@ -43,7 +100,7 @@ class CompileGraph {
     /// Returns the dependency path_ids for a given path_id (called lazily on first compile).
     using resolve_fn = std::function<llvm::SmallVector<std::uint32_t>(std::uint32_t path_id)>;
 
-    CompileGraph(dispatch_fn dispatch, resolve_fn resolve);
+    CompileGraph(kota::event_loop& loop, dispatch_fn dispatch, resolve_fn resolve);
 
     /// Compile a unit and all its transitive dependencies.
     kota::task<bool> compile(std::uint32_t path_id);
@@ -53,33 +110,95 @@ class CompileGraph {
     kota::task<bool> compile_deps(std::uint32_t path_id);
 
     /// Mark path_id and all transitive dependents as dirty,
-    /// cancelling any in-progress compilations.
+    /// cancelling any in-progress compilations (their results are stale).
     /// Returns the set of all path_ids that were marked dirty.
     llvm::SmallVector<std::uint32_t> update(std::uint32_t path_id);
 
+    /// Cancel every in-flight round regardless of interest. Waiters that
+    /// still hold interest respawn their units afterwards.
     void cancel_all();
 
+    /// Cancel all unit tasks and wait for their frames to unwind. Must be
+    /// awaited exactly once before the graph is destroyed; no compilation
+    /// can be started afterwards.
+    kota::task<> shutdown();
+
     bool has_unit(std::uint32_t path_id) const;
     bool is_dirty(std::uint32_t path_id) const;
     bool is_compiling(std::uint32_t path_id) const;
 
+    /// Current in-flight interest count for a unit (testing/diagnostics).
+    std::uint32_t refcount(std::uint32_t path_id) const;
+
+    /// All bookkeeping is quiesced: nothing compiling, no interest held and
+    /// every round's completion has fired. Holds whenever no request is in
+    /// flight and all unit tasks have unwound (e.g. after shutdown()).
+    bool idle() const;
+
+    /// Structural sanity that holds at every drain boundary: a compiling
+    /// unit has an unfinished round, and a finished round never leaves the
+    /// compiling flag behind.
+    bool consistent() const;
+
 private:
+    struct RefGuard;
+    struct UnitGuard;
+
     /// Get or create a unit, resolving its dependencies if needed.
     void ensure_resolved(std::uint32_t path_id);
 
-    /// Internal compile with ancestor tracking for cycle detection.
-    kota::task<bool> compile_impl(std::uint32_t path_id,
-                                  llvm::DenseSet<std::uint32_t> ancestors,
-                                  bool dispatch_self = true);
-
-    /// Check if waiting on `target` would deadlock given our `ancestors` chain.
-    /// Walks the dependency graph through compiling units to see if any dep
-    /// transitively reaches a unit in our ancestor chain.
-    bool has_wait_cycle(std::uint32_t target, const llvm::DenseSet<std::uint32_t>& ancestors) const;
+    /// Interest +1; creates the unit if needed.
+    void acquire(std::uint32_t path_id);
+
+    /// Interest -1; schedules a zero-interest cancellation check when it
+    /// drops to zero mid-compile.
+    void release(std::uint32_t path_id);
+
+    /// Cancels path_id's round if its interest is still zero one event-loop
+    /// tick after release() saw it drop. The delay lets transient drops
+    /// survive: a retry respawning the unit after update() re-acquires its
+    /// retained dependencies within the same drain cycle, so their in-flight
+    /// compilations are handed over instead of being killed and restarted.
+    kota::task<> zero_interest_check(std::uint32_t path_id);
+
+    /// Cancel the current round and install a fresh cancellation scope.
+    void cancel_round(CompileUnit& unit);
+
+    /// Start a unit task for path_id in the graph's task group.
+    /// Returns false when the graph is shutting down.
+    bool spawn_unit(std::uint32_t path_id);
+
+    /// One compilation round of a unit, cancelled only through its round token.
+    kota::task<> unit_body(std::uint32_t path_id, std::shared_ptr<CompileUnit::Round> round);
+
+    /// Spawned wrapper: runs unit_body under the round token and absorbs the
+    /// cancellation outcome. The task group treats a cancelled child as a
+    /// reason to abort every sibling (structured fail-fast); a round
+    /// cancelled via its token is a normal event and must not do that.
+    kota::task<> unit_task(std::uint32_t path_id,
+                           std::shared_ptr<CompileUnit::Round> round,
+                           kota::cancellation_token token);
+
+    /// Wait until path_id reaches a terminal outcome, respawning its unit
+    /// task whenever a stale round invalidates the previous attempt.
+    /// `waiter` is the unit doing the waiting (for deadlock detection),
+    /// or nullopt for requests.
+    kota::task<bool> await_unit(std::uint32_t path_id, std::optional<std::uint32_t> waiter);
+
+    /// Check if waiting on `target` would deadlock: walks the dependency
+    /// graph through compiling units to see if any dependency transitively
+    /// reaches the waiting unit.
+    bool has_wait_cycle(std::uint32_t target, std::uint32_t waiter) const;
 
     dispatch_fn dispatch;
     resolve_fn resolve;
     llvm::DenseMap<std::uint32_t, CompileUnit> units;
+
+    /// Owns every unit task; structured shutdown via shutdown().
+    /// Note: kota::task_group only reclaims completed child frames on
+    /// destruction, so frames accumulate over the graph's lifetime — one per
+    /// compilation round, same trade-off as Compiler::compile_tasks.
+    kota::task_group<> tasks;
 };
 
 }  // namespace clice

src/server/compiler/compiler.cpp

@@ -40,6 +40,12 @@ Compiler::~Compiler() {
 kota::task<> Compiler::stop() {
     compile_tasks.cancel();
     co_await compile_tasks.join();
+
+    // Requests have unwound and released their interest; now tear down the
+    // module compile graph's own unit tasks.
+    if(workspace.compile_graph) {
+        co_await workspace.compile_graph->shutdown();
+    }
 }
 
 void Compiler::init_compile_graph() {
@@ -147,7 +153,7 @@ void Compiler::init_compile_graph() {
     };
 
     workspace.compile_graph =
-        std::make_unique<CompileGraph>(std::move(dispatch), std::move(resolve));
+        std::make_unique<CompileGraph>(loop, std::move(dispatch), std::move(resolve));
     LOG_INFO("CompileGraph initialized with {} module(s)", workspace.path_to_module.size());
 }
 
@@ -562,11 +568,23 @@ kota::task<bool> Compiler::ensure_deps(Session& session,
                                        const std::string& directory,
                                        const std::vector<std::string>& arguments,
                                        std::pair<std::string, uint32_t>& pch,
-                                       std::unordered_map<std::string, std::string>& pcms) {
+                                       std::unordered_map<std::string, std::string>& pcms,
+                                       std::optional<kota::cancellation_token> scope) {
     auto path_id = session.path_id;
 
+    // Compile module dependencies within the request scope: cancelling the
+    // scope unwinds the wait and releases this request's interest in the
+    // dependency graph, without touching the shared compilations themselves.
+    auto compile_deps = [&](std::uint32_t pid) -> kota::task<bool> {
+        if(!scope) {
+            co_return co_await workspace.compile_graph->compile_deps(pid);
+        }
+        auto result = co_await kota::with_token(workspace.compile_graph->compile_deps(pid), *scope);
+        co_return result.has_value() && *result;
+    };
+
     // Compile C++20 module dependencies (PCMs).
-    if(workspace.compile_graph && !co_await workspace.compile_graph->compile_deps(path_id)) {
+    if(workspace.compile_graph && !co_await compile_deps(path_id)) {
         co_return false;
     }
 
@@ -584,7 +602,7 @@ kota::task<bool> Compiler::ensure_deps(Session& session,
                     // If PCM not already built, try to build it.
                     if(workspace.pcm_paths.find(pid) == workspace.pcm_paths.end()) {
                         if(workspace.compile_graph && workspace.compile_graph->has_unit(pid)) {
-                            co_await workspace.compile_graph->compile_deps(pid);
+                            co_await compile_deps(pid);
                         }
                     }
                     found = true;
@@ -597,6 +615,12 @@ kota::task<bool> Compiler::ensure_deps(Session& session,
         }
     }
 
+    // The buffer-scan waits above tolerate failed PCM builds, but a cancelled
+    // scope means this round was superseded — abandon it before the PCH step.
+    if(scope && scope->cancelled()) {
+        co_return false;
+    }
+
     // Build or reuse PCH.
     auto pch_ok = co_await ensure_pch(session, directory, arguments);
     if(pch_ok) {
@@ -673,7 +697,14 @@ kota::task<> Compiler::run_compile(std::uint32_t pid, std::shared_ptr<Session::P
         co_return;
     }
 
-    if(!co_await ensure_deps(*sess, params.directory, params.arguments, params.pch, params.pcms)) {
+    bool deps_ok = co_await ensure_deps(*sess,
+                                        params.directory,
+                                        params.arguments,
+                                        params.pch,
+                                        params.pcms,
+                                        pc->deps_scope.token());
+    pc->deps_done = true;
+    if(!deps_ok) {
         LOG_WARN("Dependency preparation failed for {}, skipping compile", uri_str);
         finish_compile();
         co_return;
@@ -685,6 +716,18 @@ kota::task<> Compiler::run_compile(std::uint32_t pid, std::shared_ptr<Session::P
         co_return;
     }
 
+    // Superseded while preparing dependencies — don't send the stale text:
+    // the replacement compile may already have sent newer text, and the
+    // worker applies compiles in arrival order without a version check.
+    if(sess->generation != gen) {
+        LOG_INFO("ensure_compiled: superseded before send ({} vs {}) for {}",
+                 sess->generation,
+                 gen,
+                 uri_str);
+        finish_compile();
+        co_return;
+    }
+
     auto result = co_await pool.send_stateful(pid, params);
 
     sess = find_session();
@@ -770,24 +813,42 @@ kota::task<bool> Compiler::ensure_compiled(Session& session) {
         session.ast_dirty = true;
     }
 
-    // If another compile is already in flight, wait for it.
+    // If an up-to-date compile is already in flight, wait for it.
     // This co_await may be cancelled by LSP $/cancelRequest — that's fine,
     // it just means this particular feature request is abandoned.  The
     // detached compile task keeps running independently.
     while(session.compiling) {
         auto pending = session.compiling;
+        if(pending->generation != session.generation && !pending->deps_done) {
+            // The in-flight compile is stale (user edited since it started)
+            // and still holds interest in the module graph — supersede it.
+            // A stale compile already past its dependency phase is left to
+            // finish instead: superseding it gains nothing (the worker send
+            // is not cancellable), and waiting coalesces rapid edits into a
+            // single follow-up compile at the latest generation.
+            break;
+        }
         co_await pending->done.wait();
         if(!session.ast_dirty)
             co_return true;
     }
 
+    auto superseded = session.compiling;
     auto pending_compile = std::make_shared<Session::PendingCompile>();
+    pending_compile->generation = session.generation;
     session.compiling = pending_compile;
 
     LOG_INFO("ensure_compiled: launching compile path_id={} gen={}", path_id, session.generation);
 
+    // Spawn the replacement before cancelling the superseded compile: the new
+    // round acquires its module-dependency interest synchronously, so shared
+    // dependencies never see their interest drop to zero across the swap.
     compile_tasks.spawn(run_compile(path_id, pending_compile));
 
+    if(superseded) {
+        superseded->deps_scope.cancel();
+    }
+
     // Wait for the detached compile to finish.  If this wait is cancelled
     // by LSP $/cancelRequest, the detached task continues unaffected.
     co_await pending_compile->done.wait();

src/server/compiler/compiler.h

@@ -109,11 +109,14 @@ class Compiler {
 private:
     kota::task<> run_compile(std::uint32_t path_id, std::shared_ptr<Session::PendingCompile> pc);
 
+    /// @param scope  When set, cancels the module-dependency wait if this
+    ///               compile round is superseded by a newer one.
     kota::task<bool> ensure_deps(Session& session,
                                  const std::string& directory,
                                  const std::vector<std::string>& arguments,
                                  std::pair<std::string, uint32_t>& pch,
-                                 std::unordered_map<std::string, std::string>& pcms);
+                                 std::unordered_map<std::string, std::string>& pcms,
+                                 std::optional<kota::cancellation_token> scope = {});
 
     kota::task<bool> ensure_pch(Session& session,
                                 const std::string& directory,

src/server/compiler/indexer.cpp

@@ -16,6 +16,7 @@
 #include "kota/ipc/lsp/position.h"
 #include "kota/ipc/lsp/protocol.h"
 #include "kota/ipc/lsp/uri.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"