[clangd] Export preamble AST and serialized size as metrics

[adamcz]

clang-tools-extra/clangd/Preamble.cpp

500

This is a significant change. You are now exporting this information for failed preamble builds. Do you think these are interesting? It might make the data a bit confusing (i.e. lots of small "built in <1s, did not read any files) when something is really wrong with preamble.

I think we should either only export on successful preamble OR add a boolean "success" dimension to the metric.

[sammccall]

clang-tools-extra/clangd/Preamble.cpp

500

This doesn't change what is exported, that's still controlled by TUScheduler in the same way (only exporting on success).

It just changes the behavior of this function: it now populates stats instead of leaving them uninitialized. The function already returns a success dimension (either the returned pointer is null or not).

clang-tools-extra/clangd/Preamble.h

 /// Timings and statistics from the premble build. Unlike PreambleData, these
 /// do not need to be stored for later, but can be useful for logging, metrics,
 /// etc.
 struct PreambleBuildStats {
   /// Total wall time it took to build preamble, in seconds.
   double TotalBuildTime;
   /// Time spent in filesystem operations during the build, in seconds.
   double FileSystemTime;
+
+  /// Estimate of the memory used while building the preamble.
+  /// This memory has been released when buildPreamble returns.
+  /// For example, this includes the size of the in-memory AST (ASTContext).
+  size_t BuildSize;
+  /// The serialized size of the preamble.
+  /// This storage is needed while the preamble is used (but may be on disk).
+  size_t SerializedSize;
 };
 
 /// Build a preamble for the new inputs unless an old one can be reused.
 /// If \p PreambleCallback is set, it will be run on top of the AST while
 /// building the preamble.
 /// If Stats is not non-null, build statistics will be exported there.
 std::shared_ptr<const PreambleData>
 buildPreamble(PathRef FileName, CompilerInvocation CI,

clang-tools-extra/clangd/Preamble.cpp

bool compileCommandsAreEqual(const tooling::CompileCommand &LHS,
                              const tooling::CompileCommand &RHS) {
   // We don't check for Output, it should not matter to clangd.
   return LHS.Directory == RHS.Directory && LHS.Filename == RHS.Filename &&
          llvm::makeArrayRef(LHS.CommandLine).equals(RHS.CommandLine);
 }
 
 class CppFilePreambleCallbacks : public PreambleCallbacks {
 public:
-  CppFilePreambleCallbacks(PathRef File, PreambleParsedCallback ParsedCallback)
-      : File(File), ParsedCallback(ParsedCallback) {}
+  CppFilePreambleCallbacks(PathRef File, PreambleParsedCallback ParsedCallback,
+                           PreambleBuildStats *Stats)
+      : File(File), ParsedCallback(ParsedCallback), Stats(Stats) {}
 
   IncludeStructure takeIncludes() { return std::move(Includes); }
 
   MainFileMacros takeMacros() { return std::move(Macros); }
 
   std::vector<PragmaMark> takeMarks() { return std::move(Marks); }
 
   CanonicalIncludes takeCanonicalIncludes() { return std::move(CanonIncludes); }
 
   bool isMainFileIncludeGuarded() const { return IsMainFileIncludeGuarded; }
 
   void AfterExecute(CompilerInstance &CI) override {
     if (ParsedCallback) {
       trace::Span Tracer("Running PreambleCallback");
       ParsedCallback(CI.getASTContext(), CI.getPreprocessor(), CanonIncludes);
     }
 
     const SourceManager &SM = CI.getSourceManager();
     const FileEntry *MainFE = SM.getFileEntryForID(SM.getMainFileID());
     IsMainFileIncludeGuarded =
         CI.getPreprocessor().getHeaderSearchInfo().isFileMultipleIncludeGuarded(
             MainFE);
+
+    if (Stats) {
+      const ASTContext &AST = CI.getASTContext();
+      Stats->BuildSize = AST.getASTAllocatedMemory();
+      Stats->BuildSize += AST.getSideTableAllocatedMemory();
+      Stats->BuildSize += AST.Idents.getAllocator().getTotalMemory();
+      Stats->BuildSize += AST.Selectors.getTotalMemory();
+
+      Stats->BuildSize += AST.getSourceManager().getContentCacheSize();
+      Stats->BuildSize += AST.getSourceManager().getDataStructureSizes();
+      Stats->BuildSize +=
+          AST.getSourceManager().getMemoryBufferSizes().malloc_bytes;
+
+      const Preprocessor &PP = CI.getPreprocessor();
+      Stats->BuildSize += PP.getTotalMemory();
+      if (PreprocessingRecord *PRec = PP.getPreprocessingRecord())
+        Stats->BuildSize += PRec->getTotalMemory();
+      Stats->BuildSize += PP.getHeaderSearchInfo().getTotalMemory();
+    }
   }
 
   void BeforeExecute(CompilerInstance &CI) override {
     CanonIncludes.addSystemHeadersMapping(CI.getLangOpts());
     LangOpts = &CI.getLangOpts();
     SourceMgr = &CI.getSourceManager();
     Includes.collect(CI);
   }
private:
   IncludeStructure Includes;
   CanonicalIncludes CanonIncludes;
   MainFileMacros Macros;
   std::vector<PragmaMark> Marks;
   bool IsMainFileIncludeGuarded = false;
   std::unique_ptr<CommentHandler> IWYUHandler = nullptr;
   const clang::LangOptions *LangOpts = nullptr;
   const SourceManager *SourceMgr = nullptr;
+  PreambleBuildStats *Stats;
 };
 
 // Represents directives other than includes, where basic textual information is
 // enough.
 struct TextualPPDirective {
   unsigned DirectiveLine;
   // Full text that's representing the directive, including the `#`.
   std::string Text;
buildPreamble(PathRef FileName, CompilerInvocation CI,
   // Skip function bodies when building the preamble to speed up building
   // the preamble and make it smaller.
   assert(!CI.getFrontendOpts().SkipFunctionBodies);
   CI.getFrontendOpts().SkipFunctionBodies = true;
   // We don't want to write comment locations into PCH. They are racy and slow
   // to read back. We rely on dynamic index for the comments instead.
   CI.getPreprocessorOpts().WriteCommentListToPCH = false;
 
-  CppFilePreambleCallbacks CapturedInfo(FileName, PreambleCallback);
+  CppFilePreambleCallbacks CapturedInfo(FileName, PreambleCallback, Stats);
   auto VFS = Inputs.TFS->view(Inputs.CompileCommand.Directory);
   llvm::SmallString<32> AbsFileName(FileName);
   VFS->makeAbsolute(AbsFileName);
   auto StatCache = std::make_unique<PreambleFileStatusCache>(AbsFileName);
   auto StatCacheFS = StatCache->getProducingFS(VFS);
   llvm::IntrusiveRefCntPtr<TimerFS> TimedFS(new TimerFS(StatCacheFS));
 
   WallTimer PreambleTimer;
   PreambleTimer.startTimer();
   auto BuiltPreamble = PrecompiledPreamble::Build(
       CI, ContentsBuffer.get(), Bounds, *PreambleDiagsEngine,
       Stats ? TimedFS : StatCacheFS, std::make_shared<PCHContainerOperations>(),
       StoreInMemory, CapturedInfo);
   PreambleTimer.stopTimer();
 
   // When building the AST for the main file, we do want the function
   // bodies.
   CI.getFrontendOpts().SkipFunctionBodies = false;
 
+  if (Stats != nullptr) {
+    Stats->TotalBuildTime = PreambleTimer.getTime();
+    Stats->FileSystemTime = TimedFS->getTime();
+    Stats->SerializedSize = BuiltPreamble ? BuiltPreamble->getSize() : 0;
+  }
+
   if (BuiltPreamble) {
     vlog("Built preamble of size {0} for file {1} version {2} in {3} seconds",
          BuiltPreamble->getSize(), FileName, Inputs.Version,
          PreambleTimer.getTime());
     std::vector<Diag> Diags = PreambleDiagnostics.take();
     auto Result = std::make_shared<PreambleData>(std::move(*BuiltPreamble));
     Result->Version = Inputs.Version;
     Result->CompileCommand = Inputs.CompileCommand;
     Result->Diags = std::move(Diags);
     Result->Includes = CapturedInfo.takeIncludes();
     Result->Macros = CapturedInfo.takeMacros();
     Result->Marks = CapturedInfo.takeMarks();
     Result->CanonIncludes = CapturedInfo.takeCanonicalIncludes();
     Result->StatCache = std::move(StatCache);
     Result->MainIsIncludeGuarded = CapturedInfo.isMainFileIncludeGuarded();
-    if (Stats != nullptr) {
-      Stats->TotalBuildTime = PreambleTimer.getTime();
-      Stats->FileSystemTime = TimedFS->getTime();
-    }
     return Result;
   }
 
   elog("Could not build a preamble for file {0} version {1}: {2}", FileName,
        Inputs.Version, BuiltPreamble.getError().message());
   return nullptr;
 }
 

clang-tools-extra/clangd/TUScheduler.cpp

 constexpr trace::Metric
     PreambleBuildFilesystemLatency("preamble_fs_latency",
                                    trace::Metric::Distribution, "build_type");
 // Tracks latency of FS operations done during a preamble build as a ratio of
 // preamble build time. build_type is same as above.
 constexpr trace::Metric PreambleBuildFilesystemLatencyRatio(
     "preamble_fs_latency_ratio", trace::Metric::Distribution, "build_type");
 
+constexpr trace::Metric PreambleBuildSize("preamble_build_size",
+                                          trace::Metric::Distribution);
+constexpr trace::Metric PreambleSerializedSize("preamble_serialized_size",
+                                               trace::Metric::Distribution);
+
 void reportPreambleBuild(const PreambleBuildStats &Stats,
                          bool IsFirstPreamble) {
   auto RecordWithLabel = [&Stats](llvm::StringRef Label) {
     PreambleBuildFilesystemLatency.record(Stats.FileSystemTime, Label);
     if (Stats.TotalBuildTime > 0) // Avoid division by zero.
       PreambleBuildFilesystemLatencyRatio.record(
           Stats.FileSystemTime / Stats.TotalBuildTime, Label);
   };
 
   static llvm::once_flag OnceFlag;
   llvm::call_once(OnceFlag, [&] { RecordWithLabel("first_build"); });
   RecordWithLabel(IsFirstPreamble ? "first_build_for_file" : "rebuild");
+
+  PreambleBuildSize.record(Stats.BuildSize);
+  PreambleSerializedSize.record(Stats.SerializedSize);
 }
 
 class ASTWorker;
 } // namespace
 
 static clang::clangd::Key<std::string> FileBeingProcessed;
 
 llvm::Optional<llvm::StringRef> TUScheduler::getFileBeingProcessedInContext() {