[clangd] Merge symbols in global-sym-builder on the go

[ilya-biryukov]

Summary

This avoids storing intermediate symbols in memory, most of which are
duplicates.
The resulting .yaml file is ~120MB, while intermediate symbols takes
more than 20GB.

[hokein]

Thanks for digging it out!

In upstream, we use InMemoryToolResults which saves all duplicated "std::string"s into the memory, I think we could optimize InMemoryToolResults by using Arena to keep the memory low, I will try it to see whether it can reduce the memory. A bonus point is that we don't need to change any client code.

[sammccall]

Is this patch still relevant after haojian's string deduplication?

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

53

Seems reasonably likely we would actually have contention here? merging per-thread (combiner) and then globally at the end (reducer) might be the way to go (might be significantly faster). But not sure how big the impact is.

[ilya-biryukov]

In D45478#1064027, @sammccall wrote:

Is this patch still relevant after haojian's string deduplication?

Apparently it does. It has an advantage of distributing the work more evenly between the program runs.
Currently the tool reports progress based on the number of files it parsed. But then it takes a lot of time to actually merge all the symbols at the end and the progress is not reported during that time.

The perfect behavior would be to remove duplicates, not just intern them :-) That shouldn't be too hard and it probably even makes sense to have a ToolExecutor that does that.

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

53

I haven't seen any noticeable performance degradation after this patch.
There should be contention, but it doesn't seem to matter much, as parsing still takes way more time.

Still makes sense to rewrite the code in order to avoid contention, will do that.

[sammccall]

In D45478#1071983, @ilya-biryukov wrote:

In D45478#1064027, @sammccall wrote:

Is this patch still relevant after haojian's string deduplication?

Apparently it does. It has an advantage of distributing the work more evenly between the program runs.
Currently the tool reports progress based on the number of files it parsed. But then it takes a lot of time to actually merge all the symbols at the end and the progress is not reported during that time.

The perfect behavior would be to remove duplicates, not just intern them :-) That shouldn't be too hard and it probably even makes sense to have a ToolExecutor that does that.

OK, can we sync up offline with @hokein and decide what to do here? I don't think it makes sense to have both these solutions to the same issue, but I don't have a strong opinion on exactly what we do.
(Other than clearly the right solution is to have a proper MR framework <half-trolling>)

[ilya-biryukov]

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

53

Looking at the traces, the threads spend ~6% of mergeSymbols time between Lock(Mut) and for(const Sym....
Improving it might be a good idea, but from my point of view it's not too bad to commit it as is.

Doing it properly is a little complicated in the current setup, since we don't have knobs to create per-thread structs and merge them later.

[sammccall]

I think this implementation will have problems in a "real" multi-machine MR framework. The lifetime of the Merger is the whole program, with output only coming at the end.
With N workers, each will store >1/N of the symbols (more because of overlap), and the streaming nature of a MR is important to its scaling.

As you know, we rely on this code in such a framework :-)
I would be much happier if we could express these constraints/patterns in the open-source code, I think the most promising approach would be to define a more complete clang mapreduce API in Tooling.
But that's way out of scope in this patch. I'm not really sure what to suggest here, happy to talk through it more tomorrow.

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

56

consider calling this consume, and the member in SymbolIndexActionFactory etc Sink, to emphasize the data flow

[ilya-biryukov]

This landed as a different change.

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

 #include "clang/Tooling/Execution.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Support/YAMLTraits.h"
+#include <mutex>
 
 using namespace llvm;
 using namespace clang::tooling;
 using clang::clangd::SymbolSlab;
 
 namespace clang {
 namespace clangd {
 namespace {
 
 static llvm::cl::opt<std::string> AssumedHeaderDir(
     "assume-header-dir",
     llvm::cl::desc("The index includes header that a symbol is defined in. "
                    "If the absolute path cannot be determined (e.g. an "
                    "in-memory VFS) then the relative path is resolved against "
                    "this directory, which must be absolute. If this flag is "
                    "not given, such headers will have relative paths."),
     llvm::cl::init(""));
 
+/// Combines occurrences of the same symbols across translation units.
+class SymbolMerger {
+public:
+  void mergeSymbols(const SymbolSlab &Symbols) {
+    std::lock_guard<std::mutex> Lock(Mut);
+    for (const Symbol &Sym : Symbols) {
+      if (const auto *Existing = UniqueSymbols.find(Sym.ID)) {
+        Symbol::Details Scratch;
+        UniqueSymbols.insert(mergeSymbol(*Existing, Sym, &Scratch));
+      } else {
+        UniqueSymbols.insert(Sym);
+      }
+    }
+  }
+
+  SymbolSlab build() && {
+    std::lock_guard<std::mutex> Lock(Mut);
+    return std::move(UniqueSymbols).build();
+  }
+
+private:
+  std::mutex Mut;
+  SymbolSlab::Builder UniqueSymbols;
+};
+
 class SymbolIndexActionFactory : public tooling::FrontendActionFactory {
 public:
-  SymbolIndexActionFactory(tooling::ExecutionContext *Ctx) : Ctx(Ctx) {}
+  SymbolIndexActionFactory(SymbolMerger &Merger) : Merger(Merger) {}
 
   clang::FrontendAction *create() override {
     // Wraps the index action and reports collected symbols to the execution
     // context at the end of each translation unit.
     class WrappedIndexAction : public WrapperFrontendAction {
     public:
-      WrappedIndexAction(std::shared_ptr<SymbolCollector> C,
+      WrappedIndexAction(SymbolMerger &Merger,
+                         std::shared_ptr<SymbolCollector> C,
                          std::unique_ptr<CanonicalIncludes> Includes,
-                         const index::IndexingOptions &Opts,
-                         tooling::ExecutionContext *Ctx)
+                         const index::IndexingOptions &Opts)
           : WrapperFrontendAction(
                 index::createIndexingAction(C, Opts, nullptr)),
-            Ctx(Ctx), Collector(C), Includes(std::move(Includes)),
+            Merger(Merger), Collector(C), Includes(std::move(Includes)),
             PragmaHandler(collectIWYUHeaderMaps(this->Includes.get())) {}
 
       std::unique_ptr<ASTConsumer>
       CreateASTConsumer(CompilerInstance &CI, StringRef InFile) override {
         CI.getPreprocessor().addCommentHandler(PragmaHandler.get());
         return WrapperFrontendAction::CreateASTConsumer(CI, InFile);
       }
 
       void EndSourceFileAction() override {
         WrapperFrontendAction::EndSourceFileAction();
 
-        auto Symbols = Collector->takeSymbols();
-        for (const auto &Sym : Symbols) {
-          std::string IDStr;
-          llvm::raw_string_ostream OS(IDStr);
-          OS << Sym.ID;
-          Ctx->reportResult(OS.str(), SymbolToYAML(Sym));
-        }
+        Merger.mergeSymbols(Collector->takeSymbols());
       }
 
     private:
-      tooling::ExecutionContext *Ctx;
+      SymbolMerger &Merger;
       std::shared_ptr<SymbolCollector> Collector;
       std::unique_ptr<CanonicalIncludes> Includes;
       std::unique_ptr<CommentHandler> PragmaHandler;
     };
 
     index::IndexingOptions IndexOpts;
     IndexOpts.SystemSymbolFilter =
         index::IndexingOptions::SystemSymbolFilterKind::All;
     IndexOpts.IndexFunctionLocals = false;
     auto CollectorOpts = SymbolCollector::Options();
     CollectorOpts.FallbackDir = AssumedHeaderDir;
     CollectorOpts.CollectIncludePath = true;
     CollectorOpts.CountReferences = true;
     auto Includes = llvm::make_unique<CanonicalIncludes>();
     addSystemHeadersMapping(Includes.get());
     CollectorOpts.Includes = Includes.get();
     return new WrappedIndexAction(
-        std::make_shared<SymbolCollector>(std::move(CollectorOpts)),
-        std::move(Includes), IndexOpts, Ctx);
+        Merger, std::make_shared<SymbolCollector>(std::move(CollectorOpts)),
+        std::move(Includes), IndexOpts);
   }
 
-  tooling::ExecutionContext *Ctx;
+  SymbolMerger &Merger;
 };
 
-// Combine occurrences of the same symbol across translation units.
-SymbolSlab mergeSymbols(tooling::ToolResults *Results) {
-  SymbolSlab::Builder UniqueSymbols;
-  llvm::BumpPtrAllocator Arena;
-  Symbol::Details Scratch;
-  Results->forEachResult([&](llvm::StringRef Key, llvm::StringRef Value) {
-    Arena.Reset();
-    llvm::yaml::Input Yin(Value, &Arena);
-    auto Sym = clang::clangd::SymbolFromYAML(Yin, Arena);
-    clang::clangd::SymbolID ID;
-    Key >> ID;
-    if (const auto *Existing = UniqueSymbols.find(ID))
-      UniqueSymbols.insert(mergeSymbol(*Existing, Sym, &Scratch));
-    else
-      UniqueSymbols.insert(Sym);
-  });
-  return std::move(UniqueSymbols).build();
-}
-
 } // namespace
 } // namespace clangd
 } // namespace clang
 
 int main(int argc, const char **argv) {
   llvm::sys::PrintStackTraceOnErrorSignal(argv[0]);
 
   const char* Overview =
int main(int argc, const char **argv) {
   }
 
   if (!clang::clangd::AssumedHeaderDir.empty() &&
       !llvm::sys::path::is_absolute(clang::clangd::AssumedHeaderDir)) {
     llvm::errs() << "--assume-header-dir must be an absolute path.\n";
     return 1;
   }
 
-  // Map phase: emit symbols found in each translation unit.
+  clang::clangd::SymbolMerger Merger;
+  // Emit and merge symbols found in each translation unit.
   auto Err = Executor->get()->execute(
-      llvm::make_unique<clang::clangd::SymbolIndexActionFactory>(
-          Executor->get()->getExecutionContext()));
-  if (Err) {
+      llvm::make_unique<clang::clangd::SymbolIndexActionFactory>(Merger));
+  if (Err)
     llvm::errs() << llvm::toString(std::move(Err)) << "\n";
-  }
-
-  // Reduce phase: combine symbols using the ID as a key.
-  auto UniqueSymbols =
-      clang::clangd::mergeSymbols(Executor->get()->getToolResults());
-
+  // Get resulting symbols.
+  auto UniqueSymbols = std::move(Merger).build();
   // Output phase: emit YAML for result symbols.
   SymbolsToYAML(UniqueSymbols, llvm::outs());
   return 0;
 }