[clang-tools-extra] [clangd] Explicitly block until async task completes (PR #130077)

[kadir çetinkaya via cfe-commits]

https://github.com/kadircet created https://github.com/llvm/llvm-project/pull/130077

We started seeing some use-after-frees starting with https://github.com/llvm/llvm-project/pull/125433.

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


>From 9242c82981c46bd1c19b65c9be3a142cdd605d9b Mon Sep 17 00:00:00 2001
From: Kadir Cetinkaya <kadircet at google.com>
Date: Thu, 6 Mar 2025 12:20:44 +0100
Subject: [PATCH] [clangd] Explicitly block until async task completes

---
 clang-tools-extra/clangd/ClangdServer.cpp | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

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.