[clangd] Allow to merge symbols on-the-fly in global-symbol-builder

[ilya-biryukov]

Summary

The new mode avoids serializing and deserializing YAML.
This results in better performance and less memory usage. Reduce phase
is now almost instant.

The default is to use the old mode going through YAML serialization to
allow migrating MapReduce clients that require the old mode to operate
properly. After we migrate the clients, we can switch the default to
the new mode.

[ilya-biryukov]

Collecting the timings to build the index for LLVM now. Will publish when they're ready

[ioeric]

Overall looks good to me.

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

61

AFAIK, this flag would basically depend on the executor being used. If executor can provide information like whether all tasks share memory, we can make the decision totally based on the selected executor (i.e. one fewer option).

[ilya-biryukov]

• New mode

./bin/global-symbol-builder -merge-on-the-fly -executor=all-TUs > /dev/null

40079.41s user 1874.40s system 5340% cpu 13:05.56 total

• Old mode

Still running, expect it to be more than 40minutes... Will update when I have the numbers.

[ilya-biryukov]

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

61

That SG. So we could add a virtual method to executor that would provide us with this information and then we can get rid of that option, does that LG?
However, do we have non-single-binary executor implementations upstream? We need to leave a way to run the code that actually uses YAML serialization to make sure we can experiment when things break.

[ilya-biryukov]

The numbers are finally there
New mode

./bin/global-symbol-builder -merge-on-the-fly -executor=all-TUs > /dev/null
40079.41s user 1874.40s system 5340% cpu 13:05.56 total

Old mode

./bin/global-symbol-builder -merge-on-the-fly=false -executor=all-TUs  > /dev/null
39789.70s user 2014.45s system 1385% cpu 50:17.37 total

So it's almost 5 times faster now: 50 min to 13 min.

[ioeric]

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

61

Sounds good.

However, do we have non-single-binary executor implementations upstream? We need to leave a way to run the code that actually uses YAML serialization to make sure we can experiment when things break.

The framework allows this, so I'm not sure if there is other downstream implementations that do this (probably not?). But I don't think we can get rid of the YAML serialization at this point because we still dump everything to yaml in the end?

[ilya-biryukov]

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

61

I meant the intermediate YAML serialization-deserialization parts.

I'd keep an option to force using the consumer that accumulates YAML symbols, and only allow to use on-the-fly merge when the tool executor reports that it is single-process. Would allow to:

• Avoid running on-the-fly merge when the executor does not support it (multi-binary MRs, etc)
• Allow to use on-the-fly by default for executors that run in the single binary.
• Allow to run with intermediate serialization even in a single binary if -merge-on-the-fly is specified.

Does that LG?

[ioeric]

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

61

IIUC, there will basically be a flag that allows us to force using yaml intermediate serialization, even when the executor supports on-the-fly merging. By default, we use on-the-fly merging whenever possible. If that's the case, SGTM.

[ilya-biryukov]

• Don't merge on-the-fly for multi-process executors

[ilya-biryukov]

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

61

The flag is still called -merge-on-the-fly, but it's true by default and can be set to false to reenable intermediate YAML serialization

[ioeric]

lgtm

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

60

Maybe make this hidden?

[ilya-biryukov]

• Make the option hidden

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

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(""));
 
+static llvm::cl::opt<bool> MergeOnTheFly(
+    "merge-on-the-fly",
+    llvm::cl::desc(
+        "Merges symbols for each processed translation unit as soon "
+        "they become available. This results in a smaller memory "
+        "usage and an almost instant reduce stage. Optimal for running as a "
+        "standalone tool, but cannot be used with multi-process executors like "
+        "MapReduce."),
+    llvm::cl::init(true));
+
+/// Responsible for aggregating symbols from each processed file and producing
+/// the final results. All methods in this class must be thread-safe,
+/// 'consumeSymbols' may be called from multiple threads.
+class SymbolsConsumer {
+public:
+  virtual ~SymbolsConsumer() = default;
+
+  /// Consume a SymbolSlab build for a file.
+  virtual void consumeSymbols(SymbolSlab Symbols) = 0;
+  /// Produce a resulting symbol slab, by combining  occurrences of the same
+  /// symbols across translation units.
+  virtual SymbolSlab mergeResults() = 0;
+};
+
 class SymbolIndexActionFactory : public tooling::FrontendActionFactory {
 public:
-  SymbolIndexActionFactory(tooling::ExecutionContext *Ctx) : Ctx(Ctx) {}
+  SymbolIndexActionFactory(SymbolsConsumer &Consumer) : Consumer(Consumer) {}
 
   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,
                          std::unique_ptr<CanonicalIncludes> Includes,
                          const index::IndexingOptions &Opts,
-                         tooling::ExecutionContext *Ctx)
+                         SymbolsConsumer &Consumer)
           : WrapperFrontendAction(
                 index::createIndexingAction(C, Opts, nullptr)),
-            Ctx(Ctx), Collector(C), Includes(std::move(Includes)),
+            Consumer(Consumer), 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);
       }
 
public:
         if (CI.hasDiagnostics() &&
             CI.getDiagnostics().hasUncompilableErrorOccurred()) {
           llvm::errs()
               << "Found uncompilable errors in the translation unit. Igoring "
                  "collected symbols...\n";
           return;
         }
 
-        auto Symbols = Collector->takeSymbols();
-        for (const auto &Sym : Symbols) {
-          Ctx->reportResult(Sym.ID.str(), SymbolToYAML(Sym));
-        }
+        Consumer.consumeSymbols(Collector->takeSymbols());
       }
 
     private:
-      tooling::ExecutionContext *Ctx;
+      SymbolsConsumer &Consumer;
       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;
     CollectorOpts.Origin = SymbolOrigin::Static;
     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);
+        std::move(Includes), IndexOpts, Consumer);
   }
 
-  tooling::ExecutionContext *Ctx;
+  SymbolsConsumer &Consumer;
 };
 
-// Combine occurrences of the same symbol across translation units.
-SymbolSlab mergeSymbols(tooling::ToolResults *Results) {
+/// Stashes per-file results inside ExecutionContext, merges all of them at the
+/// end. Useful for running on MapReduce infrastructure to avoid keeping symbols
+/// from multiple files in memory.
+class ToolExecutorConsumer : public SymbolsConsumer {
+public:
+  ToolExecutorConsumer(ToolExecutor &Executor) : Executor(Executor) {}
+
+  void consumeSymbols(SymbolSlab Symbols) override {
+    for (const auto &Sym : Symbols)
+      Executor.getExecutionContext()->reportResult(Sym.ID.str(),
+                                                   SymbolToYAML(Sym));
+  }
+
+  SymbolSlab mergeResults() override {
     SymbolSlab::Builder UniqueSymbols;
     llvm::BumpPtrAllocator Arena;
     Symbol::Details Scratch;
-  Results->forEachResult([&](llvm::StringRef Key, llvm::StringRef Value) {
+    Executor.getToolResults()->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();
   }
 
+private:
+  ToolExecutor &Executor;
+};
+
+/// Merges symbols for each translation unit as soon as the file is processed.
+/// Optimal choice for standalone tools.
+class OnTheFlyConsumer : public SymbolsConsumer {
+public:
+  void consumeSymbols(SymbolSlab Symbols) override {
+    std::lock_guard<std::mutex> Lock(Mut);
+    for (auto &&Sym : Symbols) {
+      Symbol::Details Scratch;
+      if (const auto *Existing = Result.find(Sym.ID))
+        Result.insert(mergeSymbol(*Existing, Sym, &Scratch));
+      else
+        Result.insert(Sym);
+    }
+  }
+
+  SymbolSlab mergeResults() override {
+    std::lock_guard<std::mutex> Lock(Mut);
+    return std::move(Result).build();
+  }
+
+private:
+  std::mutex Mut;
+  SymbolSlab::Builder Result;
+};
+
 } // namespace
 } // namespace clangd
 } // namespace clang
 
 int main(int argc, const char **argv) {
   llvm::sys::PrintStackTraceOnErrorSignal(argv[0]);
 
   const char *Overview = R"(
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;
   }
 
+  if (clang::clangd::MergeOnTheFly && !Executor->get()->isSingleProcess()) {
+    llvm::errs()
+        << "Found multi-process executor, forcing the use of intermediate YAML "
+           "serialization instead of the on-the-fly merge.\n";
+    clang::clangd::MergeOnTheFly = false;
+  }
+
+  std::unique_ptr<clang::clangd::SymbolsConsumer> Consumer;
+  if (clang::clangd::MergeOnTheFly)
+    Consumer = llvm::make_unique<clang::clangd::OnTheFlyConsumer>();
+  else
+    Consumer =
+        llvm::make_unique<clang::clangd::ToolExecutorConsumer>(**Executor);
+
   // Map phase: emit symbols found in each translation unit.
   auto Err = Executor->get()->execute(
-      llvm::make_unique<clang::clangd::SymbolIndexActionFactory>(
-          Executor->get()->getExecutionContext()));
+      llvm::make_unique<clang::clangd::SymbolIndexActionFactory>(*Consumer));
   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());
-
+  // Reduce phase: combine symbols with the same IDs.
+  auto UniqueSymbols = Consumer->mergeResults();
   // Output phase: emit YAML for result symbols.
   SymbolsToYAML(UniqueSymbols, llvm::outs());
   return 0;
 }