[clangd] Explicitly block until async task completes

[kadircet]

We started seeing some use-after-frees starting with #125433.

This patch ensures we explicitly block for the async task, if there's one,
before destructing std::future, independent of the stdlib implementation.

[llvmbot]

@llvm/pr-subscribers-clang-tools-extra

@llvm/pr-subscribers-clangd

Author: kadir çetinkaya (kadircet)

Changes

We started seeing some use-after-frees starting with #125433.

This patch ensures we explicitly block for the async task, if there's one,
before destructing std::future, independent of the stdlib implementation.

---

Full diff: https://github.com/llvm/llvm-project/pull/130077.diff

1 Files Affected:

• (modified) clang-tools-extra/clangd/ClangdServer.cpp (+11-5)

diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index 52be15d3da936..49a97da2bfa42 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -460,18 +460,24 @@ void ClangdServer::codeComplete(PathRef File, Position Pos,
     CodeCompleteResult Result = clangd::codeComplete(
         File, Pos, IP->Preamble, ParseInput, CodeCompleteOpts,
         SpecFuzzyFind ? &*SpecFuzzyFind : nullptr);
+    // We don't want `codeComplete` to wait for the async call if it doesn't use
+    // the result (e.g. non-index completion, speculation fails), so that `CB`
+    // is called as soon as results are available.
     {
       clang::clangd::trace::Span Tracer("Completion results callback");
       CB(std::move(Result));
     }
-    if (SpecFuzzyFind && SpecFuzzyFind->NewReq) {
+    if (!SpecFuzzyFind)
+      return;
+    if (SpecFuzzyFind->NewReq) {
       std::lock_guard<std::mutex> Lock(CachedCompletionFuzzyFindRequestMutex);
       CachedCompletionFuzzyFindRequestByFile[File] = *SpecFuzzyFind->NewReq;
     }
-    // SpecFuzzyFind is only destroyed after speculative fuzzy find finishes.
-    // We don't want `codeComplete` to wait for the async call if it doesn't use
-    // the result (e.g. non-index completion, speculation fails), so that `CB`
-    // is called as soon as results are available.
+    // Explicitly block until async task completes, this is fine as we've
+    // already provided reply to the client and running as a preamble task
+    // (i.e. won't block other preamble tasks).
+    if (SpecFuzzyFind->Result.valid())
+      SpecFuzzyFind->Result.wait();
   };
 
   // We use a potentially-stale preamble because latency is critical here.

[ilya-biryukov]

LGTM.

PS it would've probably not been my first choice to fix this at the use site rather than at the declaration of the type itself. I am worried that someone else using this elsewhere may end up running into the same problem.
However, I don't have strong opinions on it, so approving anyway. Just wanted to hear the rationale you have in mind for doing it at call site rather than defining a destructor that waits.

[kadircet]

However, I don't have strong opinions on it, so approving anyway. Just wanted to hear the rationale you have in mind for doing it at call site rather than defining a destructor that waits.

I didn't want to change the default destruction behavior from stdlib in places that doesn't care.

[ilya-biryukov]

I didn't want to change the default destruction behavior from stdlib in places that doesn't care.

Makes sense, keeps the struct simple. Now that we don't promise that wait() in the comment it's even better.
Thanks for clarifying!