[clangd] Cleanup: make the diags callback global in TUScheduler

[ilya-biryukov]

The tests still model the old callbacks using callbacks, unfortunately I see no good way to test same things in any other way.

[sammccall]

Nice!
The biggest suggestion I have is merging the callback into ASTCallbacks. That's awkward in initialization (makeUpdateCallbacks probably needs to be called unconditionally with a maybe-null pointer) but it seems like a natural grouping in all the places it gets plumbed around to... and there are a lot!

clangd/ClangdServer.h

232

I'm not sure using context here buys much: there aren't many layers, they're fairly coupled, and this information would look pretty natural in the interfaces.

What about:

• move the definition of DocVersion to TUScheduler
• make DocVersion a member of ParseInputs
• pass (PathRef, DocVersion, vector<Diag>) to the TUScheduler's diag callback

clangd/TUScheduler.h

93

ISTM the callback would fit nicely on the ASTCallbacks?
It even gets plumbed through to ASTWorker correctly by reference.

unittests/clangd/TUSchedulerTests.cpp

32

(up to you if you feel like fixing while here)
this is just llvm::consumeError

36

Yeah, these are messy. I also don't see something better in general.

There's a place or two below where I think we can avoid them without much verbosity, which I think is probably worthwhile even if lots of tests need them :-/

38

ParseInputs is always getInputs(File), you could do that here (similarly for updateWithDiags) to avoid repetition.

501

can we just increment a counter in the callback, and have DoUpdate look for a delta?

[ilya-biryukov]

• Merge the diagnostics callback into ParsingCallbacks
• updateWithDiags(File, getInputs(File...)) -> updateWithDiags(File, ...)

clangd/ClangdServer.h

232

I'd keep it separate: ParseInputs is defined in ClangdUnit.cpp and it serves the purpose of defining everything we need to run the compiler.
Adding DocVersion there would make no sense: it's of no use to the actual code doing preamble builds or parsing.

I would rather aim to get rid of DocVersion (for this particular purpose, there are other ways to fix the raciness that the DocVersions workaround).
WDYT?

clangd/TUScheduler.h

93

Done. This is definitely less plumbing, but IMO the public interface looks a bit worse now.
Consuming ASTs and diagnostics are two independent concerns and it's a bit awkward to combine them in the same struct.

But we don't have too many callers, updating those is easy.

unittests/clangd/TUSchedulerTests.cpp

32

Will fix in a separate commit, like it more when cleanups are the only thing in the commit.

38

It's actually a bit more complicated: getInputs(File, Contents), note the second parameter.
Done anyway, definitely looks better (and less prone to mistakes!).

The calls are now updateWith(S, File, Contents...), there's one place that still relies on the Inputs though (checks the VFS does not change in the following read).

501

Maybe... I haven't looked closely into the test, would prefer to keep this change mostly mechanical and avoid digging into the test logic.
Happy to take a look into simplifying the test separately, of course.

[sammccall]

Just the context thing. Rest is some optional simplifications.

clangd/ClangdServer.cpp

77

hmm, the double-indirection for diags seems a bit funny, especially since we have a dedicated method on ClangdServer so the lambda is trivial.

Maybe we should just make the CB class nested or friend in ClangdServer, and pass it a ClangdServer*? Then it can access DynamicIdx and consumeDiagnostics directly. The decoupling here isn't that strong to start with.

clangd/ClangdServer.h

232

Yeah, separating it from ParseInputs makes sense conceptually, and getting rid of it is indeed better.

I don't think this justifies using Context.

Passing it as a separate parameter seems fine to me. Putting it in ParseInputs with a FIXME is a little less principled but avoids interface churn.

clangd/TUScheduler.cpp

156

nit: we generally have the private section at the top without explicitly introducing it (as common in LLVM), or at the bottom. I'm fine with either style but would prefer not to add a third.

clangd/TUScheduler.h

93

(Separate yes, but I'm not sure they're that much more separate than consuming preambles vs ASTs, really. Outside of the 2/1 split of where the data goes today...)

unittests/clangd/TUSchedulerTests.cpp

38

Exactly, thanks!
(I forgot to mention Contents in the comment, doh)

501

It seems like a trivial change, and it'd be nice not to depend on the new machinery you're introducing here when we don't have to.
I'm not going to insist, but this is just...

atomic<unsigned> DiagsReceived;
TUScheduler S(..., [&](Diags) { ++ DiagsReceived }, ...);
...
DoUpdate = [&] -> bool { // as before
  unsigned PrevDiags = DiagsReceived;
  S.update(...);  blockUntilIdle(); // as before
  return DiagsReceived != PrevDiags;
}

[ilya-biryukov]

• Remove accidentally added redundant 'private:' section

[ilya-biryukov]

clangd/ClangdServer.cpp

77

I would avoid doing that, because ClangdServer is not thread-safe and giving the async callbacks that run on a different threads an access to the full class is scary.
Giving a narrower interface seems better, even if that means double-indirection

After D54829, there's no need for double-indirection too, the callback is simply directly calling into the diagnostics consumer, so it shouldn't be an issue.

clangd/ClangdServer.h

232

The versions are not too hard to remove, here's the change that does this: D54829
That would avoid the need for putting them into either the context or the ParseInputs.

clangd/TUScheduler.cpp

156

Yeah, sorry, that was an accidental change.

clangd/TUScheduler.h

93

but I'm not sure they're that much more separate than consuming preambles vs ASTs

ASTs and resulting diagnostics are on different layers of abstraction, so I don't think they fit well in the same interface, even if the data-flow for those is similar.

[ilya-biryukov]

Would like to land D54829 first and rebase on top of it.

[ilya-biryukov]

• Rebase
• Remove makeUpdateCallbacks, expose the callback class instead

[sammccall]

Everything looks so much better...

clangd/ClangdServer.h

 private:
   /// FIXME: This stats several files to find a .clang-format file. I/O can be
   /// slow. Think of a way to cache this.
   llvm::Expected<tooling::Replacements>
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
   typedef uint64_t DocVersion;
+  /// A unique version of the document stashed in the context on calls to
+  /// addDocument. Used to avoid races when sending diagnostics to the clients.
+  static Key<ClangdServer::DocVersion> DocVersionKey;
 
   void consumeDiagnostics(PathRef File, DocVersion Version,
                           std::vector<Diag> Diags);
 
   tooling::CompileCommand getCompileCommand(PathRef File);
 
   const GlobalCompilationDatabase &CDB;
   DiagnosticsConsumer &DiagConsumer;

clangd/ClangdServer.cpp

void handle(tooling::AtomicChanges SourceReplacements) override {
     Result = std::move(SourceReplacements);
   }
 
   Optional<Expected<tooling::AtomicChanges>> Result;
 };
 } // namespace
 
 // Returns callbacks that can be used to update the FileIndex with new ASTs.
-static std::unique_ptr<ParsingCallbacks>
-makeUpdateCallbacks(FileIndex *FIndex) {
+static std::unique_ptr<ParsingCallbacks> makeUpdateCallbacks(
+    FileIndex *FIndex,
+    llvm::unique_function<void(PathRef, std::vector<Diag>)> OnDiags) {
+  using DiagsCallback = decltype(OnDiags);
   struct CB : public ParsingCallbacks {
-    CB(FileIndex *FIndex) : FIndex(FIndex) {}
-    FileIndex *FIndex;
+    CB(FileIndex *FIndex, DiagsCallback OnDiags)
+        : FIndex(FIndex), OnDiags(std::move(OnDiags)) {}
 
     void onPreambleAST(PathRef Path, ASTContext &Ctx,
                        std::shared_ptr<clang::Preprocessor> PP) override {
+      if (FIndex)
         FIndex->updatePreamble(Path, Ctx, std::move(PP));
     }
 
     void onMainAST(PathRef Path, ParsedAST &AST) override {
+      if (FIndex)
         FIndex->updateMain(Path, AST);
     }
+
+    void onDiagnostics(PathRef File, std::vector<Diag> Diags) override {
+      OnDiags(File, std::move(Diags));
+    }
+
+  private:
+    FileIndex *FIndex;
+    DiagsCallback OnDiags;
   };
-  return llvm::make_unique<CB>(FIndex);
+  return llvm::make_unique<CB>(FIndex, std::move(OnDiags));
 }
 
 ClangdServer::Options ClangdServer::optsForTest() {
   ClangdServer::Options Opts;
   Opts.UpdateDebounce = std::chrono::steady_clock::duration::zero(); // Faster!
   Opts.StorePreamblesInMemory = true;
   Opts.AsyncThreadsCount = 4; // Consistent!
   return Opts;
: CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
       WorkspaceRoot(Opts.WorkspaceRoot),
       PCHs(std::make_shared<PCHContainerOperations>()),
       // Pass a callback into `WorkScheduler` to extract symbols from a newly
       // parsed file and rebuild the file index synchronously each time an AST
       // is parsed.
       // FIXME(ioeric): this can be slow and we may be able to index on less
       // critical paths.
       WorkScheduler(Opts.AsyncThreadsCount, Opts.StorePreamblesInMemory,
-                    DynamicIdx ? makeUpdateCallbacks(DynamicIdx.get())
-                               : nullptr,
+                    makeUpdateCallbacks(
+                        DynamicIdx.get(),
+                        [this](PathRef File, std::vector<Diag> Diags) {
+                          DocVersion DocVer =
+                              Context::current().getExisting(DocVersionKey);
+                          this->consumeDiagnostics(File, DocVer,
+                                                   std::move(Diags));
+                        }),
                     Opts.UpdateDebounce, Opts.RetentionPolicy) {
   if (DynamicIdx && Opts.StaticIndex) {
     MergedIdx =
         llvm::make_unique<MergedIndex>(DynamicIdx.get(), Opts.StaticIndex);
     Index = MergedIdx.get();
   } else if (DynamicIdx)
     Index = DynamicIdx.get();
   else if (Opts.StaticIndex)
     Index = Opts.StaticIndex;
   else
     Index = nullptr;
 }
 
 void ClangdServer::addDocument(PathRef File, StringRef Contents,
                                WantDiagnostics WantDiags) {
   DocVersion Version = ++InternalVersion[File];
+  WithContextValue Ctx(DocVersionKey, Version);
   ParseInputs Inputs = {getCompileCommand(File), FSProvider.getFileSystem(),
                         Contents.str()};
-
   Path FileStr = File.str();
-  WorkScheduler.update(File, std::move(Inputs), WantDiags,
-                       [this, FileStr, Version](std::vector<Diag> Diags) {
-                         consumeDiagnostics(FileStr, Version, std::move(Diags));
-                       });
+  WorkScheduler.update(File, std::move(Inputs), WantDiags);
 }
 
 void ClangdServer::removeDocument(PathRef File) {
   ++InternalVersion[File];
   WorkScheduler.remove(File);
 }
 
 void ClangdServer::codeComplete(PathRef File, Position Pos,
auto Action = [Pos](Callback<Optional<Hover>> CB,
     if (!InpAST)
       return CB(InpAST.takeError());
     CB(clangd::getHover(InpAST->AST, Pos));
   };
 
   WorkScheduler.runWithAST("Hover", File, Bind(Action, std::move(CB)));
 }
 
+Key<ClangdServer::DocVersion> ClangdServer::DocVersionKey;
+
 void ClangdServer::consumeDiagnostics(PathRef File, DocVersion Version,
                                       std::vector<Diag> Diags) {
   // We need to serialize access to resulting diagnostics to avoid calling
   // `onDiagnosticsReady` in the wrong order.
   std::lock_guard<std::mutex> DiagsLock(DiagnosticsMutex);
   DocVersion &LastReportedDiagsVersion = ReportedDiagnosticVersions[File];
 
   // FIXME(ibiryukov): get rid of '<' comparison here. In the current

clangd/TUScheduler.h

public:
   /// The \p AST always contains all AST nodes for the main file itself, and
   /// only a portion of the AST nodes deserialized from the preamble. Note that
   /// some nodes from the preamble may have been deserialized and may also be
   /// accessed from the main file AST, e.g. redecls of functions from preamble,
   /// etc. Clients are expected to process only the AST nodes from the main file
   /// in this callback (obtained via ParsedAST::getLocalTopLevelDecls) to obtain
   /// optimal performance.
   virtual void onMainAST(PathRef Path, ParsedAST &AST) {}
+
+  /// Called whenever the diagnostics for \p File are produced.
+  virtual void onDiagnostics(PathRef File, std::vector<Diag> Diags) {}
 };
 
 /// Handles running tasks for ClangdServer and managing the resources (e.g.,
 /// preambles and ASTs) for opened files.
 /// TUScheduler is not thread-safe, only one thread should be providing updates
 /// and scheduling tasks.
 /// Callbacks are run on a threadpool and it's appropriate to do slow work in
 /// them. Each task has a name, used for tracing (should be UpperCamelCase).
 /// FIXME(sammccall): pull out a scheduler options struct.
 class TUScheduler {
 public:
   TUScheduler(unsigned AsyncThreadsCount, bool StorePreamblesInMemory,
               std::unique_ptr<ParsingCallbacks> ASTCallbacks,
               std::chrono::steady_clock::duration UpdateDebounce,
               ASTRetentionPolicy RetentionPolicy);
   ~TUScheduler();
 
   /// Returns estimated memory usage for each of the currently open files.
   /// The order of results is unspecified.
   std::vector<std::pair<Path, std::size_t>> getUsedBytesPerFile() const;
 
   /// Returns a list of files with ASTs currently stored in memory. This method
   /// is not very reliable and is only used for test. E.g., the results will not
   /// contain files that currently run something over their AST.
   std::vector<Path> getFilesWithCachedAST() const;
 
   /// Schedule an update for \p File. Adds \p File to a list of tracked files if
   /// \p File was not part of it before.
-  /// FIXME(ibiryukov): remove the callback from this function.
-  void update(PathRef File, ParseInputs Inputs, WantDiagnostics WD,
-              llvm::unique_function<void(std::vector<Diag>)> OnUpdated);
+  void update(PathRef File, ParseInputs Inputs, WantDiagnostics WD);
 
   /// Remove \p File from the list of tracked files and schedule removal of its
   /// resources.
   void remove(PathRef File);
 
   /// Schedule an async task with no dependencies.
   void run(llvm::StringRef Name, llvm::unique_function<void()> Action);
 

clangd/TUScheduler.cpp

 /// dedicated thread.
 /// The ASTWorker that manages the AST is shared by both the processing thread
 /// and the TUScheduler. The TUScheduler should discard an ASTWorker when
 /// remove() is called, but its thread may be busy and we don't want to block.
 /// So the workers are accessed via an ASTWorkerHandle. Destroying the handle
 /// signals the worker to exit its run loop and gives up shared ownership of the
 /// worker.
 class ASTWorker {
+private:
   friend class ASTWorkerHandle;
   ASTWorker(PathRef FileName, TUScheduler::ASTCache &LRUCache,
             Semaphore &Barrier, bool RunSync,
             steady_clock::duration UpdateDebounce,
             std::shared_ptr<PCHContainerOperations> PCHs,
             bool StorePreamblesInMemory, ParsingCallbacks &Callbacks);
 
 public:
   /// Create a new ASTWorker and return a handle to it.
   /// The processing thread is spawned using \p Tasks. However, when \p Tasks
   /// is null, all requests will be processed on the calling thread
   /// synchronously instead. \p Barrier is acquired when processing each
   /// request, it is used to limit the number of actively running threads.
-  static ASTWorkerHandle create(PathRef FileName,
-                                TUScheduler::ASTCache &IdleASTs,
-                                AsyncTaskRunner *Tasks, Semaphore &Barrier,
-                                steady_clock::duration UpdateDebounce,
-                                std::shared_ptr<PCHContainerOperations> PCHs,
-                                bool StorePreamblesInMemory,
-                                ParsingCallbacks &Callbacks);
+  static ASTWorkerHandle create(
+      PathRef FileName, TUScheduler::ASTCache &IdleASTs, AsyncTaskRunner *Tasks,
+      Semaphore &Barrier, steady_clock::duration UpdateDebounce,
+      std::shared_ptr<PCHContainerOperations> PCHs, bool StorePreamblesInMemory,
   ~ASTWorker();
 
-  void update(ParseInputs Inputs, WantDiagnostics,
-              llvm::unique_function<void(std::vector<Diag>)> OnUpdated);
+  void update(ParseInputs Inputs, WantDiagnostics);
   void runWithAST(StringRef Name,
                   unique_function<void(Expected<InputsAndAST>)> Action);
   bool blockUntilIdle(Deadline Timeout) const;
 
   std::shared_ptr<const PreambleData> getPossiblyStalePreamble() const;
   /// Obtain a preamble reflecting all updates so far. Threadsafe.
   /// It may be delivered immediately, or later on the worker thread.
   void getCurrentPreamble(
private:
   TUScheduler::ASTCache &IdleASTs;
   const bool RunSync;
   /// Time to wait after an update to see whether another update obsoletes it.
   const steady_clock::duration UpdateDebounce;
   /// File that ASTWorker is responsible for.
   const Path FileName;
   /// Whether to keep the built preambles in memory or on disk.
   const bool StorePreambleInMemory;
-  /// Callback, invoked when preamble or main file AST is built.
+  /// Callback invoked when preamble or main file AST is built.
   ParsingCallbacks &Callbacks;
   /// Helper class required to build the ASTs.
   const std::shared_ptr<PCHContainerOperations> PCHs;
 
   Semaphore &Barrier;
   /// Inputs, corresponding to the current state of AST.
   ParseInputs FileInputs;
   /// Whether the diagnostics for the current FileInputs were reported to the
ASTWorker::~ASTWorker() {
   IdleASTs.take(this);
 #ifndef NDEBUG
   std::lock_guard<std::mutex> Lock(Mutex);
   assert(Done && "handle was not destroyed");
   assert(Requests.empty() && "unprocessed requests when destroying ASTWorker");
 #endif
 }
 
-void ASTWorker::update(ParseInputs Inputs, WantDiagnostics WantDiags,
-                       unique_function<void(std::vector<Diag>)> OnUpdated) {
-  auto Task = [=](decltype(OnUpdated) OnUpdated) mutable {
+void ASTWorker::update(ParseInputs Inputs, WantDiagnostics WantDiags) {
+  auto Task = [=]() mutable {
     // Will be used to check if we can avoid rebuilding the AST.
     bool InputsAreTheSame =
         std::tie(FileInputs.CompileCommand, FileInputs.Contents) ==
         std::tie(Inputs.CompileCommand, Inputs.Contents);
 
     tooling::CompileCommand OldCommand = std::move(FileInputs.CompileCommand);
     bool PrevDiagsWereReported = DiagsWereReported;
     FileInputs = Inputs;
if (!AST) {
           buildAST(FileName, std::move(Invocation), Inputs, NewPreamble, PCHs);
       AST = NewAST ? llvm::make_unique<ParsedAST>(std::move(*NewAST)) : nullptr;
     }
     // We want to report the diagnostics even if this update was cancelled.
     // It seems more useful than making the clients wait indefinitely if they
     // spam us with updates.
     // Note *AST can still be null if buildAST fails.
     if (*AST) {
-      OnUpdated((*AST)->getDiagnostics());
+      Callbacks.onDiagnostics(FileName, (*AST)->getDiagnostics());
       trace::Span Span("Running main AST callback");
       Callbacks.onMainAST(FileName, **AST);
       DiagsWereReported = true;
     }
     // Stash the AST in the cache for further use.
     IdleASTs.put(this, std::move(*AST));
   };
 
-  startTask("Update", Bind(Task, std::move(OnUpdated)), WantDiags);
+  startTask("Update", std::move(Task), WantDiags);
 }
 
 void ASTWorker::runWithAST(
     StringRef Name, unique_function<void(Expected<InputsAndAST>)> Action) {
   auto Task = [=](decltype(Action) Action) {
     Optional<std::unique_ptr<ParsedAST>> AST = IdleASTs.take(this);
     if (!AST) {
       std::unique_ptr<CompilerInvocation> Invocation =
if (!File.getValue()->Worker->blockUntilIdle(D))
       return false;
   if (PreambleTasks)
     if (!PreambleTasks->wait(D))
       return false;
   return true;
 }
 
 void TUScheduler::update(PathRef File, ParseInputs Inputs,
-                         WantDiagnostics WantDiags,
-                         unique_function<void(std::vector<Diag>)> OnUpdated) {
+                         WantDiagnostics WantDiags) {
   std::unique_ptr<FileData> &FD = Files[File];
   if (!FD) {
     // Create a new worker to process the AST-related tasks.
     ASTWorkerHandle Worker = ASTWorker::create(
         File, *IdleASTs, WorkerThreads ? WorkerThreads.getPointer() : nullptr,
         Barrier, UpdateDebounce, PCHOps, StorePreamblesInMemory, *Callbacks);
     FD = std::unique_ptr<FileData>(new FileData{
         Inputs.Contents, Inputs.CompileCommand, std::move(Worker)});
   } else {
     FD->Contents = Inputs.Contents;
     FD->Command = Inputs.CompileCommand;
   }
-  FD->Worker->update(std::move(Inputs), WantDiags, std::move(OnUpdated));
+  FD->Worker->update(std::move(Inputs), WantDiags);
 }
 
 void TUScheduler::remove(PathRef File) {
   bool Removed = Files.erase(File);
   if (!Removed)
     elog("Trying to remove file from TUScheduler that is not tracked: {0}",
          File);
 }

unittests/clangd/TUSchedulerTests.cpp

 //===-- TUSchedulerTests.cpp ------------------------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Context.h"
 #include "TUScheduler.h"
 #include "TestFS.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <algorithm>
 #include <utility>
 
 using namespace llvm;
 namespace clang {
 namespace clangd {
 namespace {
 
 using ::testing::_;
 using ::testing::AnyOf;
 using ::testing::Each;
 using ::testing::ElementsAre;
 using ::testing::Pair;
 using ::testing::Pointee;
 using ::testing::UnorderedElementsAre;
 
-void ignoreUpdate(Optional<std::vector<Diag>>) {}
 void ignoreError(Error Err) {
   handleAllErrors(std::move(Err), [](const ErrorInfoBase &) {});
 }
 
 class TUSchedulerTests : public ::testing::Test {
 protected:
   ParseInputs getInputs(PathRef File, std::string Contents) {
     return ParseInputs{*CDB.getCompileCommand(File),
                        buildTestFS(Files, Timestamps), std::move(Contents)};
   }
 
+void updateWithCallback(TUScheduler &S, PathRef File, StringRef Contents,
+                        WantDiagnostics WD, llvm::unique_function<void()> CB) {
+  WithContextValue Ctx(llvm::make_scope_exit(std::move(CB)));
+  S.update(File, getInputs(File, Contents), WD);
+}
+
+static Key<llvm::unique_function<void(PathRef File, std::vector<Diag>)>>
+    DiagsCallbackKey;
+
+/// A diagnostics callback that should be passed to TUScheduler when it's used
+/// in updateWithDiags.
+static std::unique_ptr<ParsingCallbacks> captureDiags() {
+  class CaptureDiags : public ParsingCallbacks {
+    void onDiagnostics(PathRef File, std::vector<Diag> Diags) override {
+      auto D = Context::current().get(DiagsCallbackKey);
+      if (!D)
+        return;
+      const_cast<llvm::unique_function<void(PathRef, std::vector<Diag>)> &> (
+          *D)(File, Diags);
+    }
+  };
+  return llvm::make_unique<CaptureDiags>();
+}
+
+/// Schedule an update and call \p CB with the diagnostics it produces, if any.
+/// The TUScheduler should be created with captureDiags as a DiagsCallback for
+/// this to work.
+void updateWithDiags(TUScheduler &S, PathRef File, ParseInputs Inputs,
+                     WantDiagnostics WD,
+                     llvm::unique_function<void(std::vector<Diag>)> CB) {
+  Path OrigFile = File.str();
+  WithContextValue Ctx(
+      DiagsCallbackKey,
+      Bind(
+          [OrigFile](decltype(CB) CB, PathRef File, std::vector<Diag> Diags) {
+            assert(File == OrigFile);
+            CB(std::move(Diags));
+          },
+          std::move(CB)));
+  S.update(File, std::move(Inputs), WD);
+}
+
+  void updateWithDiags(TUScheduler &S, PathRef File, llvm::StringRef Contents,
+                      WantDiagnostics WD,
+                      llvm::unique_function<void(std::vector<Diag>)> CB) {
+    return updateWithDiags(S, File, getInputs(File, Contents), WD, std::move(CB));
+  }
+
   StringMap<std::string> Files;
   StringMap<time_t> Timestamps;
   MockCompilationDatabase CDB;
 };
 
+Key<llvm::unique_function<void(PathRef File, std::vector<Diag>)>>
+    TUSchedulerTests::DiagsCallbackKey;
+
 TEST_F(TUSchedulerTests, MissingFiles) {
   TUScheduler S(getDefaultAsyncThreadsCount(),
                 /*StorePreamblesInMemory=*/true, /*ASTCallbacks=*/nullptr,
                 /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero(),
                 ASTRetentionPolicy());
 
   auto Added = testPath("added.cpp");
   Files[Added] = "";
 
   auto Missing = testPath("missing.cpp");
   Files[Missing] = "";
 
-  S.update(Added, getInputs(Added, ""), WantDiagnostics::No, ignoreUpdate);
+  S.update(Added, getInputs(Added, ""), WantDiagnostics::No);
 
   // Assert each operation for missing file is an error (even if it's available
   // in VFS).
   S.runWithAST("", Missing, [&](Expected<InputsAndAST> AST) {
     ASSERT_FALSE(bool(AST));
     ignoreError(AST.takeError());
   });
   S.runWithPreamble("", Missing, TUScheduler::Stale,

 TEST_F(TUSchedulerTests, WantDiagnostics) {
   std::atomic<int> CallbackCount(0);
   {
     // To avoid a racy test, don't allow tasks to actualy run on the worker
     // thread until we've scheduled them all.
     Notification Ready;
     TUScheduler S(
         getDefaultAsyncThreadsCount(),
-        /*StorePreamblesInMemory=*/true, /*ASTCallbacks=*/nullptr,
+        /*StorePreamblesInMemory=*/true, captureDiags(),
         /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero(),
         ASTRetentionPolicy());
     auto Path = testPath("foo.cpp");
-    S.update(Path, getInputs(Path, ""), WantDiagnostics::Yes,
+    updateWithDiags(S, Path, "", WantDiagnostics::Yes,
                     [&](std::vector<Diag>) { Ready.wait(); });
-
-    S.update(Path, getInputs(Path, "request diags"), WantDiagnostics::Yes,
-             [&](std::vector<Diag> Diags) { ++CallbackCount; });
-    S.update(Path, getInputs(Path, "auto (clobbered)"), WantDiagnostics::Auto,
-             [&](std::vector<Diag> Diags) {
-               ADD_FAILURE() << "auto should have been cancelled by auto";
+    updateWithDiags(S, Path, "request diags", WantDiagnostics::Yes,
+                    [&](std::vector<Diag>) { ++CallbackCount; });
+    updateWithDiags(S, Path, "auto (clobbered)",
+                    WantDiagnostics::Auto, [&](std::vector<Diag>) {
+                      ADD_FAILURE()
+                          << "auto should have been cancelled by auto";
                     });
-    S.update(Path, getInputs(Path, "request no diags"), WantDiagnostics::No,
-             [&](std::vector<Diag> Diags) {
+    updateWithDiags(S, Path, "request no diags", WantDiagnostics::No,
+                    [&](std::vector<Diag>) {
                       ADD_FAILURE() << "no diags should not be called back";
                     });
-    S.update(Path, getInputs(Path, "auto (produces)"), WantDiagnostics::Auto,
-             [&](std::vector<Diag> Diags) { ++CallbackCount; });
+    updateWithDiags(S, Path, "auto (produces)", WantDiagnostics::Auto,
+                    [&](std::vector<Diag>) { ++CallbackCount; });
     Ready.notify();
   }
   EXPECT_EQ(2, CallbackCount);
 }
 
 TEST_F(TUSchedulerTests, Debounce) {
   std::atomic<int> CallbackCount(0);
   {
     TUScheduler S(getDefaultAsyncThreadsCount(),
-                  /*StorePreamblesInMemory=*/true, /*ASTCallbacks=*/nullptr,
+                  /*StorePreamblesInMemory=*/true, captureDiags(),
                   /*UpdateDebounce=*/std::chrono::seconds(1),
                   ASTRetentionPolicy());
     // FIXME: we could probably use timeouts lower than 1 second here.
     auto Path = testPath("foo.cpp");
-    S.update(Path, getInputs(Path, "auto (debounced)"), WantDiagnostics::Auto,
-             [&](std::vector<Diag> Diags) {
-               ADD_FAILURE() << "auto should have been debounced and canceled";
+    updateWithDiags(S, Path, "auto (debounced)", WantDiagnostics::Auto,
+                    [&](std::vector<Diag>) {
+                      ADD_FAILURE()
+                          << "auto should have been debounced and canceled";
                     });
     std::this_thread::sleep_for(std::chrono::milliseconds(200));
-    S.update(Path, getInputs(Path, "auto (timed out)"), WantDiagnostics::Auto,
-             [&](std::vector<Diag> Diags) { ++CallbackCount; });
+    updateWithDiags(S, Path, "auto (timed out)",
+                    WantDiagnostics::Auto,
+                    [&](std::vector<Diag>) { ++CallbackCount; });
     std::this_thread::sleep_for(std::chrono::seconds(2));
-    S.update(Path, getInputs(Path, "auto (shut down)"), WantDiagnostics::Auto,
-             [&](std::vector<Diag> Diags) { ++CallbackCount; });
+    updateWithDiags(S, Path, "auto (shut down)", WantDiagnostics::Auto,
+                    [&](std::vector<Diag>) { ++CallbackCount; });
   }
   EXPECT_EQ(2, CallbackCount);
 }
 
 static std::vector<std::string> includes(const PreambleData *Preamble) {
   std::vector<std::string> Result;
   if (Preamble)
     for (const auto &Inclusion : Preamble->Includes.MainFileIncludes)
TUScheduler S(
         getDefaultAsyncThreadsCount(), /*StorePreamblesInMemory=*/true,
         /*ASTCallbacks=*/nullptr,
         /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero(),
         ASTRetentionPolicy());
     auto Path = testPath("foo.cpp");
     // Schedule two updates (A, B) and two preamble reads (stale, consistent).
     // The stale read should see A, and the consistent read should see B.
     // (We recognize the preambles by their included files).
-    S.update(Path, getInputs(Path, "#include <A>"), WantDiagnostics::Yes,
-             [&](std::vector<Diag> Diags) {
-               // This callback runs in between the two preamble updates.
+    updateWithCallback(
+        S, Path, "#include <A>", WantDiagnostics::Yes, [&]() {
 
-               // This blocks update B, preventing it from winning the race
-               // against the stale read.
-               // If the first read was instead consistent, this would deadlock.
-               InconsistentReadDone.wait();
-               // This delays update B, preventing it from winning a race
-               // against the consistent read. The consistent read sees B
-               // only because it waits for it.
-               // If the second read was stale, it would usually see A.
-               std::this_thread::sleep_for(std::chrono::milliseconds(100));
-             });
-    S.update(Path, getInputs(Path, "#include <B>"), WantDiagnostics::Yes,
-             [&](std::vector<Diag> Diags) {});
+    S.update(Path, getInputs(Path, "#include <B>"), WantDiagnostics::Yes);
 
     S.runWithPreamble("StaleRead", Path, TUScheduler::Stale,
                       [&](Expected<InputsAndPreamble> Pre) {
                         ASSERT_TRUE(bool(Pre));
                         assert(bool(Pre));
                         EXPECT_THAT(includes(Pre->Preamble),
                                     ElementsAre("<A>"));
                         InconsistentReadDone.notify();
TEST_F(TUSchedulerTests, ManyUpdates) {
   std::mutex Mut;
   int TotalASTReads = 0;
   int TotalPreambleReads = 0;
   int TotalUpdates = 0;
 
   // Run TUScheduler and collect some stats.
   {
     TUScheduler S(getDefaultAsyncThreadsCount(),
-                  /*StorePreamblesInMemory=*/true, /*ASTCallbacks=*/nullptr,
+                  /*StorePreamblesInMemory=*/true, captureDiags(),
                   /*UpdateDebounce=*/std::chrono::milliseconds(50),
                   ASTRetentionPolicy());
 
     std::vector<std::string> Files;
     for (int I = 0; I < FilesCount; ++I) {
       std::string Name = "foo" + std::to_string(I) + ".cpp";
       Files.push_back(testPath(Name));
       this->Files[Files.back()] = "";
// Run TUScheduler and collect some stats.
     int Nonce = 0;
 
     for (int FileI = 0; FileI < FilesCount; ++FileI) {
       for (int UpdateI = 0; UpdateI < UpdatesPerFile; ++UpdateI) {
         auto Contents = AllContents[(FileI + UpdateI) % AllContentsSize];
 
         auto File = Files[FileI];
         auto Inputs = getInputs(File, Contents.str());
-
         {
           WithContextValue WithNonce(NonceKey, ++Nonce);
-          S.update(File, Inputs, WantDiagnostics::Auto,
-                   [File, Nonce, &Mut,
-                    &TotalUpdates](Optional<std::vector<Diag>> Diags) {
-                     EXPECT_THAT(Context::current().get(NonceKey),
-                                 Pointee(Nonce));
+          updateWithDiags(
+              S, File, Inputs, WantDiagnostics::Auto,
+              [File, Nonce, &Mut, &TotalUpdates](std::vector<Diag>) {
+                EXPECT_THAT(Context::current().get(NonceKey), Pointee(Nonce));
 
-                     std::lock_guard<std::mutex> Lock(Mut);
-                     ++TotalUpdates;
-                     EXPECT_EQ(File,
-                               *TUScheduler::getFileBeingProcessedInContext());
-                   });
+                EXPECT_EQ(File, *TUScheduler::getFileBeingProcessedInContext());
         }
-
         {
           WithContextValue WithNonce(NonceKey, ++Nonce);
           S.runWithAST("CheckAST", File,
                        [File, Inputs, Nonce, &Mut,
                         &TotalASTReads](Expected<InputsAndAST> AST) {
                          EXPECT_THAT(Context::current().get(NonceKey),
                                      Pointee(Nonce));
 
TEST_F(TUSchedulerTests, EvictedAST) {
   )cpp";
 
   auto Foo = testPath("foo.cpp");
   auto Bar = testPath("bar.cpp");
   auto Baz = testPath("baz.cpp");
 
   // Build one file in advance. We will not access it later, so it will be the
   // one that the cache will evict.
-  S.update(Foo, getInputs(Foo, SourceContents), WantDiagnostics::Yes,
-           [&BuiltASTCounter](std::vector<Diag> Diags) { ++BuiltASTCounter; });
+  updateWithCallback(S, Foo, SourceContents,
+                     WantDiagnostics::Yes,
+                     [&BuiltASTCounter]() { ++BuiltASTCounter; });
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   ASSERT_EQ(BuiltASTCounter.load(), 1);
 
   // Build two more files. Since we can retain only 2 ASTs, these should be the
   // ones we see in the cache later.
-  S.update(Bar, getInputs(Bar, SourceContents), WantDiagnostics::Yes,
-           [&BuiltASTCounter](std::vector<Diag> Diags) { ++BuiltASTCounter; });
-  S.update(Baz, getInputs(Baz, SourceContents), WantDiagnostics::Yes,
-           [&BuiltASTCounter](std::vector<Diag> Diags) { ++BuiltASTCounter; });
+  updateWithCallback(S, Bar, SourceContents,
+                     WantDiagnostics::Yes,
+                     [&BuiltASTCounter]() { ++BuiltASTCounter; });
+  updateWithCallback(S, Baz, SourceContents,
+                     WantDiagnostics::Yes,
+                     [&BuiltASTCounter]() { ++BuiltASTCounter; });
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   ASSERT_EQ(BuiltASTCounter.load(), 3);
 
   // Check only the last two ASTs are retained.
   ASSERT_THAT(S.getFilesWithCachedAST(), UnorderedElementsAre(Bar, Baz));
 
   // Access the old file again.
-  S.update(Foo, getInputs(Foo, OtherSourceContents), WantDiagnostics::Yes,
-           [&BuiltASTCounter](std::vector<Diag> Diags) { ++BuiltASTCounter; });
+  updateWithCallback(S, Foo, OtherSourceContents, WantDiagnostics::Yes,
+                     [&BuiltASTCounter]() { ++BuiltASTCounter; });
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   ASSERT_EQ(BuiltASTCounter.load(), 4);
 
   // Check the AST for foo.cpp is retained now and one of the others got
   // evicted.
   EXPECT_THAT(S.getFilesWithCachedAST(),
               UnorderedElementsAre(Foo, AnyOf(Bar, Baz)));
 }
TEST_F(TUSchedulerTests, EmptyPreamble) {
 
   Files[Header] = "void foo()";
   Timestamps[Header] = time_t(0);
   auto WithPreamble = R"cpp(
     #include "foo.h"
     int main() {}
   )cpp";
   auto WithEmptyPreamble = R"cpp(int main() {})cpp";
-  S.update(Foo, getInputs(Foo, WithPreamble), WantDiagnostics::Auto,
-           [](std::vector<Diag>) {});
+  S.update(Foo, getInputs(Foo, WithPreamble), WantDiagnostics::Auto);
   S.runWithPreamble("getNonEmptyPreamble", Foo, TUScheduler::Stale,
                     [&](Expected<InputsAndPreamble> Preamble) {
                       // We expect to get a non-empty preamble.
                       EXPECT_GT(cantFail(std::move(Preamble))
                                     .Preamble->Preamble.getBounds()
                                     .Size,
                                 0u);
                     });
   // Wait for the preamble is being built.
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
 
   // Update the file which results in an empty preamble.
-  S.update(Foo, getInputs(Foo, WithEmptyPreamble), WantDiagnostics::Auto,
-           [](std::vector<Diag>) {});
+  S.update(Foo, getInputs(Foo, WithEmptyPreamble), WantDiagnostics::Auto);
   // Wait for the preamble is being built.
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   S.runWithPreamble("getEmptyPreamble", Foo, TUScheduler::Stale,
                     [&](Expected<InputsAndPreamble> Preamble) {
                       // We expect to get an empty preamble.
                       EXPECT_EQ(cantFail(std::move(Preamble))
                                     .Preamble->Preamble.getBounds()
                                     .Size,
auto NonEmptyPreamble = R"cpp(
     #define FOO 1
     #define BAR 2
 
     int main() {}
   )cpp";
   constexpr int ReadsToSchedule = 10;
   std::mutex PreamblesMut;
   std::vector<const void *> Preambles(ReadsToSchedule, nullptr);
-  S.update(Foo, getInputs(Foo, NonEmptyPreamble), WantDiagnostics::Auto,
-           [](std::vector<Diag>) {});
+  S.update(Foo, getInputs(Foo, NonEmptyPreamble), WantDiagnostics::Auto);
   for (int I = 0; I < ReadsToSchedule; ++I) {
     S.runWithPreamble(
         "test", Foo, TUScheduler::Stale,
         [I, &PreamblesMut, &Preambles](Expected<InputsAndPreamble> IP) {
           std::lock_guard<std::mutex> Lock(PreamblesMut);
           Preambles[I] = cantFail(std::move(IP)).Preamble;
         });
   }
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   // Check all actions got the same non-null preamble.
   std::lock_guard<std::mutex> Lock(PreamblesMut);
   ASSERT_NE(Preambles[0], nullptr);
   ASSERT_THAT(Preambles, Each(Preambles[0]));
 }
 
 TEST_F(TUSchedulerTests, NoopOnEmptyChanges) {
   TUScheduler S(
       /*AsyncThreadsCount=*/getDefaultAsyncThreadsCount(),
-      /*StorePreambleInMemory=*/true, /*ASTCallbacks=*/nullptr,
+      /*StorePreambleInMemory=*/true, captureDiags(),
       /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero(),
       ASTRetentionPolicy());
 
   auto Source = testPath("foo.cpp");
   auto Header = testPath("foo.h");
 
   Files[Header] = "int a;";
   Timestamps[Header] = time_t(0);
 
   auto SourceContents = R"cpp(
       #include "foo.h"
       int b = a;
     )cpp";
 
   // Return value indicates if the updated callback was received.
-  auto DoUpdate = [&](ParseInputs Inputs) -> bool {
+  auto DoUpdate = [&](std::string Contents) -> bool {
     std::atomic<bool> Updated(false);
     Updated = false;
-    S.update(Source, std::move(Inputs), WantDiagnostics::Yes,
+    updateWithDiags(S, Source, Contents, WantDiagnostics::Yes,
                     [&Updated](std::vector<Diag>) { Updated = true; });
     bool UpdateFinished = S.blockUntilIdle(timeoutSeconds(10));
     if (!UpdateFinished)
       ADD_FAILURE() << "Updated has not finished in one second. Threading bug?";
     return Updated;
   };
 
   // Test that subsequent updates with the same inputs do not cause rebuilds.
-  ASSERT_TRUE(DoUpdate(getInputs(Source, SourceContents)));
-  ASSERT_FALSE(DoUpdate(getInputs(Source, SourceContents)));
+  ASSERT_TRUE(DoUpdate(SourceContents));
+  ASSERT_FALSE(DoUpdate(SourceContents));
 
   // Update to a header should cause a rebuild, though.
   Files[Header] = time_t(1);
-  ASSERT_TRUE(DoUpdate(getInputs(Source, SourceContents)));
-  ASSERT_FALSE(DoUpdate(getInputs(Source, SourceContents)));
+  ASSERT_TRUE(DoUpdate(SourceContents));
+  ASSERT_FALSE(DoUpdate(SourceContents));
 
   // Update to the contents should cause a rebuild.
   auto OtherSourceContents = R"cpp(
       #include "foo.h"
       int c = d;
     )cpp";
-  ASSERT_TRUE(DoUpdate(getInputs(Source, OtherSourceContents)));
-  ASSERT_FALSE(DoUpdate(getInputs(Source, OtherSourceContents)));
+  ASSERT_TRUE(DoUpdate(OtherSourceContents));
+  ASSERT_FALSE(DoUpdate(OtherSourceContents));
 
   // Update to the compile commands should also cause a rebuild.
   CDB.ExtraClangFlags.push_back("-DSOMETHING");
-  ASSERT_TRUE(DoUpdate(getInputs(Source, OtherSourceContents)));
-  ASSERT_FALSE(DoUpdate(getInputs(Source, OtherSourceContents)));
+  ASSERT_TRUE(DoUpdate(OtherSourceContents));
+  ASSERT_FALSE(DoUpdate(OtherSourceContents));
 }
 
 TEST_F(TUSchedulerTests, NoChangeDiags) {
   TUScheduler S(
       /*AsyncThreadsCount=*/getDefaultAsyncThreadsCount(),
-      /*StorePreambleInMemory=*/true, /*ASTCallbacks=*/nullptr,
+      /*StorePreambleInMemory=*/true, captureDiags(),
       /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero(),
       ASTRetentionPolicy());
 
   auto FooCpp = testPath("foo.cpp");
   auto Contents = "int a; int b;";
 
-  S.update(FooCpp, getInputs(FooCpp, Contents), WantDiagnostics::No,
-           [](std::vector<Diag>) { ADD_FAILURE() << "Should not be called."; });
+  updateWithDiags(
+      S, FooCpp, Contents, WantDiagnostics::No,
   S.runWithAST("touchAST", FooCpp, [](Expected<InputsAndAST> IA) {
     // Make sure the AST was actually built.
     cantFail(std::move(IA));
   });
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
 
   // Even though the inputs didn't change and AST can be reused, we need to
   // report the diagnostics, as they were not reported previously.
   std::atomic<bool> SeenDiags(false);
-  S.update(FooCpp, getInputs(FooCpp, Contents), WantDiagnostics::Auto,
+  updateWithDiags(S, FooCpp, Contents, WantDiagnostics::Auto,
                   [&](std::vector<Diag>) { SeenDiags = true; });
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   ASSERT_TRUE(SeenDiags);
 
   // Subsequent request does not get any diagnostics callback because the same
   // diags have previously been reported and the inputs didn't change.
-  S.update(
-      FooCpp, getInputs(FooCpp, Contents), WantDiagnostics::Auto,
+  updateWithDiags(
+      S, FooCpp, Contents, WantDiagnostics::Auto,
       [&](std::vector<Diag>) { ADD_FAILURE() << "Should not be called."; });
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
 }
 
 TEST_F(TUSchedulerTests, Run) {
   TUScheduler S(/*AsyncThreadsCount=*/getDefaultAsyncThreadsCount(),
                 /*StorePreambleInMemory=*/true, /*ASTCallbacks=*/nullptr,
                 /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero(),