Index: clangd/ClangdServer.h
===================================================================
--- clangd/ClangdServer.h
+++ clangd/ClangdServer.h
@@ -226,20 +226,12 @@
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
-  typedef uint64_t DocVersion;
-
-  void consumeDiagnostics(PathRef File, DocVersion Version,
-                          std::vector<Diag> Diags);
-
   tooling::CompileCommand getCompileCommand(PathRef File);
 
   const GlobalCompilationDatabase &CDB;
   DiagnosticsConsumer &DiagConsumer;
   const FileSystemProvider &FSProvider;
 
-  /// Used to synchronize diagnostic responses for added and removed files.
-  llvm::StringMap<DocVersion> InternalVersion;
-
   Path ResourceDir;
   // The index used to look up symbols. This could be:
   //   - null (all index functionality is optional)
@@ -259,12 +251,6 @@
 
   llvm::Optional<std::string> WorkspaceRoot;
   std::shared_ptr<PCHContainerOperations> PCHs;
-  /// Used to serialize diagnostic callbacks.
-  /// FIXME(ibiryukov): get rid of an extra map and put all version counters
-  /// into CppFile.
-  std::mutex DiagnosticsMutex;
-  /// Maps from a filename to the latest version of reported diagnostics.
-  llvm::StringMap<DocVersion> ReportedDiagnosticVersions;
   // WorkScheduler has to be the last member, because its destructor has to be
   // called before all other members to stop the worker thread that references
   // ClangdServer.
Index: clangd/ClangdServer.cpp
===================================================================
--- clangd/ClangdServer.cpp
+++ clangd/ClangdServer.cpp
@@ -134,19 +134,17 @@
 
 void ClangdServer::addDocument(PathRef File, StringRef Contents,
                                WantDiagnostics WantDiags) {
-  DocVersion Version = ++InternalVersion[File];
   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));
+                       [this, FileStr](std::vector<Diag> Diags) {
+                         DiagConsumer.onDiagnosticsReady(FileStr,
+                                                         std::move(Diags));
                        });
 }
 
 void ClangdServer::removeDocument(PathRef File) {
-  ++InternalVersion[File];
   WorkScheduler.remove(File);
 }
 
@@ -445,24 +443,6 @@
   WorkScheduler.runWithAST("Hover", File, Bind(Action, std::move(CB)));
 }
 
-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
-  // implementation diagnostics will not be reported after version counters'
-  // overflow. This should not happen in practice, since DocVersion is a
-  // 64-bit unsigned integer.
-  if (Version < LastReportedDiagsVersion)
-    return;
-  LastReportedDiagsVersion = Version;
-
-  DiagConsumer.onDiagnosticsReady(File, std::move(Diags));
-}
-
 tooling::CompileCommand ClangdServer::getCompileCommand(PathRef File) {
   trace::Span Span("GetCompileCommand");
   Optional<tooling::CompileCommand> C = CDB.getCompileCommand(File);
Index: clangd/TUScheduler.cpp
===================================================================
--- clangd/TUScheduler.cpp
+++ clangd/TUScheduler.cpp
@@ -238,6 +238,9 @@
   Semaphore &Barrier;
   /// Inputs, corresponding to the current state of AST.
   ParseInputs FileInputs;
+  /// Used to indicate to the tasks running inside the worker queue that AST is
+  /// being destroyed and the AST queue is being drained.
+  std::atomic<bool> ShuttingDown{false};
   /// Whether the diagnostics for the current FileInputs were reported to the
   /// users before.
   bool DiagsWereReported = false;
@@ -401,8 +404,9 @@
       }
     }
 
-    // We only need to build the AST if diagnostics were requested.
-    if (WantDiags == WantDiagnostics::No)
+    // We only need to build the AST if diagnostics were requested and the
+    // callers are still interested in this ASTWorker (it's not shutting down).
+    if (WantDiags == WantDiagnostics::No || ShuttingDown)
       return;
 
     // Get the AST for diagnostics.
@@ -509,6 +513,7 @@
 bool ASTWorker::isASTCached() const { return IdleASTs.getUsedBytes(this) != 0; }
 
 void ASTWorker::stop() {
+  ShuttingDown = true;
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     assert(!Done && "stop() called twice");
Index: unittests/clangd/ClangdTests.cpp
===================================================================
--- unittests/clangd/ClangdTests.cpp
+++ unittests/clangd/ClangdTests.cpp
@@ -748,7 +748,8 @@
         BlockingRequests[RequestIndex]();
       }
     }
-  } // Wait for ClangdServer to shutdown before proceeding.
+    ASSERT_TRUE(Server.blockUntilIdleForTest());
+  }
 
   // Check some invariants about the state of the program.
   std::vector<FileStat> Stats = DiagConsumer.takeFileStats();
Index: unittests/clangd/TUSchedulerTests.cpp
===================================================================
--- unittests/clangd/TUSchedulerTests.cpp
+++ unittests/clangd/TUSchedulerTests.cpp
@@ -147,6 +147,8 @@
     std::this_thread::sleep_for(std::chrono::seconds(2));
     S.update(Path, getInputs(Path, "auto (shut down)"), WantDiagnostics::Auto,
              [&](std::vector<Diag> Diags) { ++CallbackCount; });
+
+    ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   }
   EXPECT_EQ(2, CallbackCount);
 }