[clangd] Move caching of compile args out of ClangdServer.

Summary:
Caching is now handled by ClangdLSPServer and hidden behind the
GlobalCompilationDatabase interface. This simplifies ClangdServer.
This change also removes the SkipCache flag from addDocument,
which is now obsolete.

No behavioral changes are intended, the clangd binary still caches the
compile commands on the first read.

Reviewers: sammccall

Reviewed By: sammccall

Subscribers: mgorny, ioeric, MaskRay, jkorous, cfe-commits

Differential Revision: https://reviews.llvm.org/D48068

12 files changed, 109 insertions, 139 deletions

diff --git a/clang-tools-extra/clangd/CMakeLists.txt b/clang-tools-extra/clangd/CMakeLists.txt
index 6ff7fc555d6..16d5c0f2285 100644
--- a/clang-tools-extra/clangd/CMakeLists.txt
+++ b/clang-tools-extra/clangd/CMakeLists.txt
@@ -14,7 +14,6 @@ add_clang_library(clangDaemon
   ClangdUnit.cpp
   CodeComplete.cpp
   CodeCompletionStrings.cpp
-  CompileArgsCache.cpp
   Compiler.cpp
   Context.cpp
   Diagnostics.cpp
diff --git a/clang-tools-extra/clangd/ClangdLSPServer.cpp b/clang-tools-extra/clangd/ClangdLSPServer.cpp
index 3e0e4f6d497..7e708a29665 100644
--- a/clang-tools-extra/clangd/ClangdLSPServer.cpp
+++ b/clang-tools-extra/clangd/ClangdLSPServer.cpp
@@ -129,11 +129,13 @@ void ClangdLSPServer::onShutdown(ShutdownParams &Params) {
 void ClangdLSPServer::onExit(ExitParams &Params) { IsDone = true; }
 
 void ClangdLSPServer::onDocumentDidOpen(DidOpenTextDocumentParams &Params) {
-  if (Params.metadata && !Params.metadata->extraFlags.empty())
-    CDB.setExtraFlagsForFile(Params.textDocument.uri.file(),
-                             std::move(Params.metadata->extraFlags));
-
   PathRef File = Params.textDocument.uri.file();
+  if (Params.metadata && !Params.metadata->extraFlags.empty()) {
+    NonCachedCDB.setExtraFlagsForFile(File,
+                                      std::move(Params.metadata->extraFlags));
+    CDB.invalidate(File);
+  }
+
   std::string &Contents = Params.textDocument.text;
 
   DraftMgr.addDraft(File, Contents);
@@ -155,6 +157,7 @@ void ClangdLSPServer::onDocumentDidChange(DidChangeTextDocumentParams &Params) {
     // fail rather than giving wrong results.
     DraftMgr.removeDraft(File);
     Server.removeDocument(File);
+    CDB.invalidate(File);
     log(llvm::toString(Contents.takeError()));
     return;
   }
@@ -383,7 +386,10 @@ void ClangdLSPServer::onChangeConfiguration(
 
   // Compilation database change.
   if (Settings.compilationDatabasePath.hasValue()) {
-    CDB.setCompileCommandsDir(Settings.compilationDatabasePath.getValue());
+    NonCachedCDB.setCompileCommandsDir(
+        Settings.compilationDatabasePath.getValue());
+    CDB.clear();
+
     reparseOpenedFiles();
   }
 }
@@ -392,8 +398,8 @@ ClangdLSPServer::ClangdLSPServer(JSONOutput &Out,
                                  const clangd::CodeCompleteOptions &CCOpts,
                                  llvm::Optional<Path> CompileCommandsDir,
                                  const ClangdServer::Options &Opts)
-    : Out(Out), CDB(std::move(CompileCommandsDir)), CCOpts(CCOpts),
-      SupportedSymbolKinds(defaultSymbolKinds()),
+    : Out(Out), NonCachedCDB(std::move(CompileCommandsDir)), CDB(NonCachedCDB),
+      CCOpts(CCOpts), SupportedSymbolKinds(defaultSymbolKinds()),
       Server(CDB, FSProvider, /*DiagConsumer=*/*this, Opts) {}
 
 bool ClangdLSPServer::run(std::FILE *In, JSONStreamStyle InputStyle) {
@@ -471,6 +477,5 @@ void ClangdLSPServer::onDiagnosticsReady(PathRef File,
 void ClangdLSPServer::reparseOpenedFiles() {
   for (const Path &FilePath : DraftMgr.getActiveFiles())
     Server.addDocument(FilePath, *DraftMgr.getDraft(FilePath),
-                       WantDiagnostics::Auto,
-                       /*SkipCache=*/true);
+                       WantDiagnostics::Auto);
 }
diff --git a/clang-tools-extra/clangd/ClangdLSPServer.h b/clang-tools-extra/clangd/ClangdLSPServer.h
index 67beed1e3cf..0a3e6fc0a52 100644
--- a/clang-tools-extra/clangd/ClangdLSPServer.h
+++ b/clang-tools-extra/clangd/ClangdLSPServer.h
@@ -100,7 +100,9 @@ private:
 
   // Various ClangdServer parameters go here. It's important they're created
   // before ClangdServer.
-  DirectoryBasedGlobalCompilationDatabase CDB;
+  DirectoryBasedGlobalCompilationDatabase NonCachedCDB;
+  CachingCompilationDb CDB;
+
   RealFileSystemProvider FSProvider;
   /// Options used for code completion
   clangd::CodeCompleteOptions CCOpts;
diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index 3eae2e8d445..c2b68f8a970 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -23,6 +23,7 @@
 #include "clang/Tooling/Refactoring/Rename/RenamingAction.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
@@ -83,9 +84,9 @@ ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
                            FileSystemProvider &FSProvider,
                            DiagnosticsConsumer &DiagConsumer,
                            const Options &Opts)
-    : CompileArgs(CDB, Opts.ResourceDir ? Opts.ResourceDir->str()
-                                        : getStandardResourceDir()),
-      DiagConsumer(DiagConsumer), FSProvider(FSProvider),
+    : CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
+      ResourceDir(Opts.ResourceDir ? Opts.ResourceDir->str()
+                                   : getStandardResourceDir()),
       FileIdx(Opts.BuildDynamicSymbolIndex ? new FileIndex() : nullptr),
       PCHs(std::make_shared<PCHContainerOperations>()),
       // Pass a callback into `WorkScheduler` to extract symbols from a newly
@@ -120,13 +121,10 @@ void ClangdServer::setRootPath(PathRef RootPath) {
 }
 
 void ClangdServer::addDocument(PathRef File, StringRef Contents,
-                               WantDiagnostics WantDiags, bool SkipCache) {
-  if (SkipCache)
-    CompileArgs.invalidate(File);
-
+                               WantDiagnostics WantDiags) {
   DocVersion Version = ++InternalVersion[File];
-  ParseInputs Inputs = {CompileArgs.getCompileCommand(File),
-                        FSProvider.getFileSystem(), Contents.str()};
+  ParseInputs Inputs = {getCompileCommand(File), FSProvider.getFileSystem(),
+                        Contents.str()};
 
   Path FileStr = File.str();
   WorkScheduler.update(File, std::move(Inputs), WantDiags,
@@ -137,7 +135,6 @@ void ClangdServer::addDocument(PathRef File, StringRef Contents,
 
 void ClangdServer::removeDocument(PathRef File) {
   ++InternalVersion[File];
-  CompileArgs.invalidate(File);
   WorkScheduler.remove(File);
 }
 
@@ -430,6 +427,17 @@ void ClangdServer::consumeDiagnostics(PathRef File, DocVersion Version,
   DiagConsumer.onDiagnosticsReady(File, std::move(Diags));
 }
 
+tooling::CompileCommand ClangdServer::getCompileCommand(PathRef File) {
+  llvm::Optional<tooling::CompileCommand> C = CDB.getCompileCommand(File);
+  if (!C) // FIXME: Suppress diagnostics? Let the user know?
+    C = CDB.getFallbackCommand(File);
+
+  // Inject the resource dir.
+  // FIXME: Don't overwrite it if it's already there.
+  C->CommandLine.push_back("-resource-dir=" + ResourceDir);
+  return std::move(*C);
+}
+
 void ClangdServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
   // FIXME: Do nothing for now. This will be used for indexing and potentially
   // invalidating other caches.
diff --git a/clang-tools-extra/clangd/ClangdServer.h b/clang-tools-extra/clangd/ClangdServer.h
index 4ccf58f4494..b6d5f160128 100644
--- a/clang-tools-extra/clangd/ClangdServer.h
+++ b/clang-tools-extra/clangd/ClangdServer.h
@@ -12,7 +12,6 @@
 
 #include "ClangdUnit.h"
 #include "CodeComplete.h"
-#include "CompileArgsCache.h"
 #include "Function.h"
 #include "GlobalCompilationDatabase.h"
 #include "Protocol.h"
@@ -122,8 +121,7 @@ public:
   /// When \p SkipCache is true, compile commands will always be requested from
   /// compilation database even if they were cached in previous invocations.
   void addDocument(PathRef File, StringRef Contents,
-                   WantDiagnostics WD = WantDiagnostics::Auto,
-                   bool SkipCache = false);
+                   WantDiagnostics WD = WantDiagnostics::Auto);
 
   /// Remove \p File from list of tracked files, schedule a request to free
   /// resources associated with it.
@@ -216,13 +214,16 @@ private:
   void consumeDiagnostics(PathRef File, DocVersion Version,
                           std::vector<Diag> Diags);
 
-  CompileArgsCache CompileArgs;
+  tooling::CompileCommand getCompileCommand(PathRef File);
+
+  GlobalCompilationDatabase &CDB;
   DiagnosticsConsumer &DiagConsumer;
   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)
   //   - the dynamic index owned by ClangdServer (FileIdx)
diff --git a/clang-tools-extra/clangd/CompileArgsCache.cpp b/clang-tools-extra/clangd/CompileArgsCache.cpp
deleted file mode 100644
index 42167e56bc3..00000000000
--- a/clang-tools-extra/clangd/CompileArgsCache.cpp
+++ /dev/null
@@ -1,44 +0,0 @@
-//===--- CompileArgsCache.cpp --------------------------------------------===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===---------------------------------------------------------------------===//
-
-#include "CompileArgsCache.h"
-
-namespace clang {
-namespace clangd {
-namespace {
-tooling::CompileCommand getCompileCommand(GlobalCompilationDatabase &CDB,
-                                          PathRef File, PathRef ResourceDir) {
-  llvm::Optional<tooling::CompileCommand> C = CDB.getCompileCommand(File);
-  if (!C) // FIXME: Suppress diagnostics? Let the user know?
-    C = CDB.getFallbackCommand(File);
-
-  // Inject the resource dir.
-  // FIXME: Don't overwrite it if it's already there.
-  C->CommandLine.push_back("-resource-dir=" + ResourceDir.str());
-  return std::move(*C);
-}
-} // namespace
-
-CompileArgsCache::CompileArgsCache(GlobalCompilationDatabase &CDB,
-                                   Path ResourceDir)
-    : CDB(CDB), ResourceDir(std::move(ResourceDir)) {}
-
-tooling::CompileCommand CompileArgsCache::getCompileCommand(PathRef File) {
-  auto It = Cached.find(File);
-  if (It == Cached.end()) {
-    It = Cached.insert({File, clangd::getCompileCommand(CDB, File, ResourceDir)})
-             .first;
-  }
-  return It->second;
-}
-
-void CompileArgsCache::invalidate(PathRef File) { Cached.erase(File); }
-
-} // namespace clangd
-} // namespace clang
diff --git a/clang-tools-extra/clangd/CompileArgsCache.h b/clang-tools-extra/clangd/CompileArgsCache.h
deleted file mode 100644
index 49d88758104..00000000000
--- a/clang-tools-extra/clangd/CompileArgsCache.h
+++ /dev/null
@@ -1,43 +0,0 @@
-//===--- CompileArgsCache.h -------------------------------------*- C++-*-===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===---------------------------------------------------------------------===//
-
-#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_COMPILEARGSCACHE_H
-#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_COMPILEARGSCACHE_H
-
-#include "GlobalCompilationDatabase.h"
-#include "Path.h"
-#include "clang/Tooling/CompilationDatabase.h"
-
-namespace clang {
-namespace clangd {
-
-/// A helper class used by ClangdServer to get compile commands from CDB.
-/// Also caches CompileCommands produced by compilation database on per-file
-/// basis. This avoids queries to CDB that can be much more expensive than a
-/// table lookup.
-class CompileArgsCache {
-public:
-  CompileArgsCache(GlobalCompilationDatabase &CDB, Path ResourceDir);
-
-  /// Gets compile command for \p File from cache or CDB if it's not in the
-  /// cache.
-  tooling::CompileCommand getCompileCommand(PathRef File);
-
-  /// Removes a cache entry for \p File, if it's present in the cache.
-  void invalidate(PathRef File);
-
-private:
-  GlobalCompilationDatabase &CDB;
-  const Path ResourceDir;
-  llvm::StringMap<tooling::CompileCommand> Cached;
-};
-
-} // namespace clangd
-} // namespace clang
-#endif
diff --git a/clang-tools-extra/clangd/GlobalCompilationDatabase.cpp b/clang-tools-extra/clangd/GlobalCompilationDatabase.cpp
index ca223f6b344..a8ac58d3d4d 100644
--- a/clang-tools-extra/clangd/GlobalCompilationDatabase.cpp
+++ b/clang-tools-extra/clangd/GlobalCompilationDatabase.cpp
@@ -119,5 +119,38 @@ DirectoryBasedGlobalCompilationDatabase::getCDBForFile(PathRef File) const {
   return nullptr;
 }
 
+CachingCompilationDb::CachingCompilationDb(
+    const GlobalCompilationDatabase &InnerCDB)
+    : InnerCDB(InnerCDB) {}
+
+llvm::Optional<tooling::CompileCommand>
+CachingCompilationDb::getCompileCommand(PathRef File) const {
+  std::unique_lock<std::mutex> Lock(Mut);
+  auto It = Cached.find(File);
+  if (It != Cached.end())
+    return It->second;
+
+  Lock.unlock();
+  llvm::Optional<tooling::CompileCommand> Command =
+      InnerCDB.getCompileCommand(File);
+  Lock.lock();
+  return Cached.try_emplace(File, std::move(Command)).first->getValue();
+}
+
+tooling::CompileCommand
+CachingCompilationDb::getFallbackCommand(PathRef File) const {
+  return InnerCDB.getFallbackCommand(File);
+}
+
+void CachingCompilationDb::invalidate(PathRef File) {
+  std::unique_lock<std::mutex> Lock(Mut);
+  Cached.erase(File);
+}
+
+void CachingCompilationDb::clear() {
+  std::unique_lock<std::mutex> Lock(Mut);
+  Cached.clear();
+}
+
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/GlobalCompilationDatabase.h b/clang-tools-extra/clangd/GlobalCompilationDatabase.h
index cf1e613fbf1..ab89a185b46 100644
--- a/clang-tools-extra/clangd/GlobalCompilationDatabase.h
+++ b/clang-tools-extra/clangd/GlobalCompilationDatabase.h
@@ -85,6 +85,34 @@ private:
   /// is located.
   llvm::Optional<Path> CompileCommandsDir;
 };
+
+/// A wrapper around GlobalCompilationDatabase that caches the compile commands.
+/// Note that only results of getCompileCommand are cached.
+class CachingCompilationDb : public GlobalCompilationDatabase {
+public:
+  explicit CachingCompilationDb(const GlobalCompilationDatabase &InnerCDB);
+
+  /// Gets compile command for \p File from cache or CDB if it's not in the
+  /// cache.
+  llvm::Optional<tooling::CompileCommand>
+  getCompileCommand(PathRef File) const override;
+
+  /// Forwards to the inner CDB. Results of this function are not cached.
+  tooling::CompileCommand getFallbackCommand(PathRef File) const override;
+
+  /// Removes an entry for \p File if it's present in the cache.
+  void invalidate(PathRef File);
+
+  /// Removes all cached compile commands.
+  void clear();
+
+private:
+  const GlobalCompilationDatabase &InnerCDB;
+  mutable std::mutex Mut;
+  mutable llvm::StringMap<llvm::Optional<tooling::CompileCommand>>
+      Cached; /* GUARDED_BY(Mut) */
+};
+
 } // namespace clangd
 } // namespace clang
 
diff --git a/clang-tools-extra/unittests/clangd/ClangdTests.cpp b/clang-tools-extra/unittests/clangd/ClangdTests.cpp
index a91b65ff2b7..ae47b60a53e 100644
--- a/clang-tools-extra/unittests/clangd/ClangdTests.cpp
+++ b/clang-tools-extra/unittests/clangd/ClangdTests.cpp
@@ -390,16 +390,7 @@ struct bar { T x; };
 
   // Now switch to C++ mode.
   CDB.ExtraClangFlags = {"-xc++"};
-  // By default addDocument does not check if CompileCommand has changed, so we
-  // expect to see the errors.
-  runAddDocument(Server, FooCpp, SourceContents1);
-  EXPECT_TRUE(DiagConsumer.hadErrorInLastDiags());
-  runAddDocument(Server, FooCpp, SourceContents2);
-  EXPECT_TRUE(DiagConsumer.hadErrorInLastDiags());
-  // Passing SkipCache=true will force addDocument to reparse the file with
-  // proper flags.
-  runAddDocument(Server, FooCpp, SourceContents2, WantDiagnostics::Auto,
-                 /*SkipCache=*/true);
+  runAddDocument(Server, FooCpp, SourceContents2, WantDiagnostics::Auto);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
   // Subsequent addDocument calls should finish without errors too.
   runAddDocument(Server, FooCpp, SourceContents1);
@@ -431,14 +422,7 @@ int main() { return 0; }
 
   // Parse without the define, no errors should be produced.
   CDB.ExtraClangFlags = {};
-  // By default addDocument does not check if CompileCommand has changed, so we
-  // expect to see the errors.
-  runAddDocument(Server, FooCpp, SourceContents);
-  EXPECT_TRUE(DiagConsumer.hadErrorInLastDiags());
-  // Passing SkipCache=true will force addDocument to reparse the file with
-  // proper flags.
-  runAddDocument(Server, FooCpp, SourceContents, WantDiagnostics::Auto,
-                 /*SkipCache=*/true);
+  runAddDocument(Server, FooCpp, SourceContents, WantDiagnostics::Auto);
   ASSERT_TRUE(Server.blockUntilIdleForTest());
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
   // Subsequent addDocument call should finish without errors too.
@@ -500,10 +484,8 @@ int hello;
   CDB.ExtraClangFlags.clear();
   DiagConsumer.clear();
   Server.removeDocument(BazCpp);
-  Server.addDocument(FooCpp, FooSource.code(), WantDiagnostics::Auto,
-                     /*SkipCache=*/true);
-  Server.addDocument(BarCpp, BarSource.code(), WantDiagnostics::Auto,
-                     /*SkipCache=*/true);
+  Server.addDocument(FooCpp, FooSource.code(), WantDiagnostics::Auto);
+  Server.addDocument(BarCpp, BarSource.code(), WantDiagnostics::Auto);
   ASSERT_TRUE(Server.blockUntilIdleForTest());
 
   EXPECT_THAT(DiagConsumer.filesWithDiags(),
@@ -708,7 +690,7 @@ int d;
       Server.addDocument(FilePaths[FileIndex],
                          ShouldHaveErrors ? SourceContentsWithErrors
                                           : SourceContentsWithoutErrors,
-                         WantDiagnostics::Auto, SkipCache);
+                         WantDiagnostics::Auto);
       UpdateStatsOnAddDocument(FileIndex, ShouldHaveErrors);
     };
 
diff --git a/clang-tools-extra/unittests/clangd/SyncAPI.cpp b/clang-tools-extra/unittests/clangd/SyncAPI.cpp
index 0a1c5988e7c..aa2a044f231 100644
--- a/clang-tools-extra/unittests/clangd/SyncAPI.cpp
+++ b/clang-tools-extra/unittests/clangd/SyncAPI.cpp
@@ -12,8 +12,8 @@ namespace clang {
 namespace clangd {
 
 void runAddDocument(ClangdServer &Server, PathRef File, StringRef Contents,
-                    WantDiagnostics WantDiags, bool SkipCache) {
-  Server.addDocument(File, Contents, WantDiags, SkipCache);
+                    WantDiagnostics WantDiags) {
+  Server.addDocument(File, Contents, WantDiags);
   if (!Server.blockUntilIdleForTest())
     llvm_unreachable("not idle after addDocument");
 }
diff --git a/clang-tools-extra/unittests/clangd/SyncAPI.h b/clang-tools-extra/unittests/clangd/SyncAPI.h
index d4d2ac8f901..0a140aa9db3 100644
--- a/clang-tools-extra/unittests/clangd/SyncAPI.h
+++ b/clang-tools-extra/unittests/clangd/SyncAPI.h
@@ -20,8 +20,7 @@ namespace clangd {
 
 // Calls addDocument and then blockUntilIdleForTest.
 void runAddDocument(ClangdServer &Server, PathRef File, StringRef Contents,
-                    WantDiagnostics WantDiags = WantDiagnostics::Auto,
-                    bool SkipCache = false);
+                    WantDiagnostics WantDiags = WantDiagnostics::Auto);
 
 llvm::Expected<CompletionList>
 runCodeComplete(ClangdServer &Server, PathRef File, Position Pos,