[clangd] Provide a helper to report estimated memory usage per-file

[ioeric]

clangd/ClangdServer.h

325

s/currently/current/ ? or current open?

327

Is the total usage of all files a good estimate of the overall clangd memory usage? It might be worth mentioning the relationship between this and overall clangd memory here.

clangd/ClangdUnit.cpp

658

Would it be possible to calculate the usages lazily on request, so that you don't need to maintain two states?

clangd/ClangdUnit.h

100

It seems that CppFile::getUsedBytes() includes preamble?

unittests/clangd/ClangdTests.cpp

428

Can we add a test for non-empty files and check that the usage is reasonable (e.g. >0)?

[ilya-biryukov]

clangd/ClangdServer.h

325

Intended this to be currently open files. Fixed, thanks!

327

It isn't.
Added a comment.

clangd/ClangdUnit.cpp

658

It's really complicated since we store std::futures. This is one of the caveats that I'm trying to fix with threading revamp, implementation will change accordingly. For now, having two states is much easier.

clangd/ClangdUnit.h

100

CppFile includes the size of the Preamble, ParsedAST doesn't

unittests/clangd/ClangdTests.cpp

428

We do exactly that.

[ilya-biryukov]

• Fixed review comments

[ioeric]

lg

clangd/ClangdUnit.cpp

658

Thanks for the explanation! A FIXME would be appropriate here then? In general, I wouldn't expect this API to be called very often.

unittests/clangd/ClangdTests.cpp

428

My bad!

struct Something {
  int method();
};

looks like real code on phabricator!

[ilya-biryukov]

• Added a FIXME to CppFile::getUsedBytes()

[ilya-biryukov]

Thanks for reviewing this!

unittests/clangd/ClangdTests.cpp

428

Phabricator is too smart :-)

clangd/ClangdServer.h

public:
 
   /// Only for testing purposes.
   /// Waits until all requests to worker thread are finished and dumps AST for
   /// \p File. \p File must be in the list of added documents.
   std::string dumpAST(PathRef File);
   /// Called when an event occurs for a watched file in the workspace.
   void onFileEvent(const DidChangeWatchedFilesParams &Params);
 
+  /// Returns estimated memory usage for each of the currently open files.
+  /// The order of results is unspecified.
+  /// Overall memory usage of clangd may be significantly more than reported
+  /// here, as this metric does not account (at least) for:
+  ///   - memory occupied by static and dynamic index,
+  ///   - memory required for in-flight requests,
+  /// FIXME: those metrics might be useful too, we should add them.
+  std::vector<std::pair<Path, std::size_t>> getUsedBytesPerFile() const;
+
 private:
   /// FIXME: This stats several files to find a .clang-format file. I/O can be
   /// slow. Think of a way to cache this.
   llvm::Expected<tooling::Replacements>
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
   std::future<Context>

clangd/ClangdServer.cpp

WorkScheduler.addToFront(std::move(CancelReparses), std::move(DonePromise),
                            std::move(DeferredCancel), std::move(Ctx));
   return DoneFuture;
 }
 
 void ClangdServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
   // FIXME: Do nothing for now. This will be used for indexing and potentially
   // invalidating other caches.
 }
+
+std::vector<std::pair<Path, std::size_t>>
+ClangdServer::getUsedBytesPerFile() const {
+  return Units.getUsedBytesPerFile();
+}

clangd/ClangdUnit.h

public:
 
   /// This function returns all top-level decls, including those that come
   /// from Preamble. Decls, coming from Preamble, have to be deserialized, so
   /// this call might be expensive.
   ArrayRef<const Decl *> getTopLevelDecls();
 
   const std::vector<DiagWithFixIts> &getDiagnostics() const;
 
+  /// Returns the esitmated size of the AST and the accessory structures, in
+  /// bytes. Does not include the size of the preamble.
+  std::size_t getUsedBytes() const;
+
 private:
   ParsedAST(std::shared_ptr<const PreambleData> Preamble,
             std::unique_ptr<CompilerInstance> Clang,
             std::unique_ptr<FrontendAction> Action,
             std::vector<const Decl *> TopLevelDecls,
             std::vector<DiagWithFixIts> Diags);
 
 private:
public:
 
   /// Returns a future to get the most fresh AST for a file. Returned AST is
   /// wrapped to prevent concurrent accesses.
   /// We use std::shared_ptr here because MVSC fails to compile non-copyable
   /// classes as template arguments of promise/future. It is guaranteed to
   /// always be non-null.
   std::shared_future<std::shared_ptr<ParsedASTWrapper>> getAST() const;
 
+  /// Returns an estimated size, in bytes, currently occupied by the AST and the
+  /// Preamble.
+  std::size_t getUsedBytes() const;
+
 private:
   /// A helper guard that manages the state of CppFile during rebuild.
   class RebuildGuard {
   public:
     RebuildGuard(CppFile &File, unsigned RequestRebuildCounter);
     ~RebuildGuard();
 
     bool wasCancelledBeforeConstruction() const;
private:
   mutable std::mutex Mutex;
   /// A counter to cancel old rebuilds.
   unsigned RebuildCounter;
   /// Used to wait when rebuild is finished before starting another one.
   bool RebuildInProgress;
   /// Condition variable to indicate changes to RebuildInProgress.
   std::condition_variable RebuildCond;
 
+  /// Size of the last built AST, in bytes.
+  std::size_t ASTMemUsage;
+  /// Size of the last build Preamble, in bytes.
+  std::size_t PreambleMemUsage;
+
   /// Promise and future for the latests AST. Fulfilled during rebuild.
   /// We use std::shared_ptr here because MVSC fails to compile non-copyable
   /// classes as template arguments of promise/future.
   std::promise<std::shared_ptr<ParsedASTWrapper>> ASTPromise;
   std::shared_future<std::shared_ptr<ParsedASTWrapper>> ASTFuture;
 
   /// Promise and future for the latests Preamble. Fulfilled during rebuild.
   std::promise<std::shared_ptr<const PreambleData>> PreamblePromise;

clangd/ClangdUnit.cpp

 #include <algorithm>
 #include <chrono>
 
 using namespace clang::clangd;
 using namespace clang;
 
 namespace {
 
+template <class T> std::size_t getUsedBytes(const std::vector<T> &Vec) {
+  return Vec.capacity() * sizeof(T);
+}
+
 class DeclTrackingASTConsumer : public ASTConsumer {
 public:
   DeclTrackingASTConsumer(std::vector<const Decl *> &TopLevelDecls)
       : TopLevelDecls(TopLevelDecls) {}
 
   bool HandleTopLevelDecl(DeclGroupRef DG) override {
     for (const Decl *D : DG) {
       // ObjCMethodDecl are not actually top-level decls.
ArrayRef<const Decl *> ParsedAST::getTopLevelDecls() {
   ensurePreambleDeclsDeserialized();
   return TopLevelDecls;
 }
 
 const std::vector<DiagWithFixIts> &ParsedAST::getDiagnostics() const {
   return Diags;
 }
 
+std::size_t ParsedAST::getUsedBytes() const {
+  auto &AST = getASTContext();
+  // FIXME(ibiryukov): we do not account for the dynamically allocated part of
+  // SmallVector<FixIt> inside each Diag.
+  return AST.getASTAllocatedMemory() + AST.getSideTableAllocatedMemory() +
+         ::getUsedBytes(TopLevelDecls) + ::getUsedBytes(Diags);
+}
+
 PreambleData::PreambleData(PrecompiledPreamble Preamble,
                            std::vector<serialization::DeclID> TopLevelDeclIDs,
                            std::vector<DiagWithFixIts> Diags)
     : Preamble(std::move(Preamble)),
       TopLevelDeclIDs(std::move(TopLevelDeclIDs)), Diags(std::move(Diags)) {}
 
 ParsedAST::ParsedAST(std::shared_ptr<const PreambleData> Preamble,
                      std::unique_ptr<CompilerInstance> Clang,
return std::shared_ptr<CppFile>(new CppFile(FileName, StorePreamblesInMemory,
                                               std::move(PCHs),
                                               std::move(ASTCallback)));
 }
 
 CppFile::CppFile(PathRef FileName, bool StorePreamblesInMemory,
                  std::shared_ptr<PCHContainerOperations> PCHs,
                  ASTParsedCallback ASTCallback)
     : FileName(FileName), StorePreamblesInMemory(StorePreamblesInMemory),
-      RebuildCounter(0), RebuildInProgress(false), PCHs(std::move(PCHs)),
+      RebuildCounter(0), RebuildInProgress(false), ASTMemUsage(0),
+      PreambleMemUsage(0), PCHs(std::move(PCHs)),
       ASTCallback(std::move(ASTCallback)) {
   // FIXME(ibiryukov): we should pass a proper Context here.
   log(Context::empty(), "Created CppFile for " + FileName);
 
   std::lock_guard<std::mutex> Lock(Mutex);
   LatestAvailablePreamble = nullptr;
   PreamblePromise.set_value(nullptr);
   PreambleFuture = PreamblePromise.get_future();
This->RebuildCond.wait(Lock, [This, RequestRebuildCounter]() {
              This->RebuildCounter != RequestRebuildCounter;
     });
 
     // This computation got cancelled itself, do nothing.
     if (This->RebuildCounter != RequestRebuildCounter)
       return;
 
     // Set empty results for Promises.
+    That->PreambleMemUsage = 0;
     That->PreamblePromise.set_value(nullptr);
+    That->ASTMemUsage = 0;
     That->ASTPromise.set_value(std::make_shared<ParsedASTWrapper>(llvm::None));
   };
 }
 
 llvm::Optional<std::vector<DiagWithFixIts>>
 CppFile::rebuild(const Context &Ctx, ParseInputs &&Inputs) {
   return deferRebuild(std::move(Inputs))(Ctx);
 }
auto FinishRebuild =
     // Publish the new Preamble.
     {
       std::lock_guard<std::mutex> Lock(That->Mutex);
       // We always set LatestAvailablePreamble to the new value, hoping that it
       // will still be usable in the further requests.
       That->LatestAvailablePreamble = NewPreamble;
       if (RequestRebuildCounter != That->RebuildCounter)
         return llvm::None; // Our rebuild request was cancelled, do nothing.
+      That->PreambleMemUsage =
+          NewPreamble ? NewPreamble->Preamble.getSize() : 0;
       That->PreamblePromise.set_value(NewPreamble);
     } // unlock Mutex
 
     // Prepare the Preamble and supplementary data for rebuilding AST.
     std::vector<DiagWithFixIts> Diagnostics;
     if (NewPreamble) {
       Diagnostics.insert(Diagnostics.begin(), NewPreamble->Diags.begin(),
                          NewPreamble->Diags.end());
auto FinishRebuild =
 
     // Publish the new AST.
     {
       std::lock_guard<std::mutex> Lock(That->Mutex);
       if (RequestRebuildCounter != That->RebuildCounter)
         return Diagnostics; // Our rebuild request was cancelled, don't set
                             // ASTPromise.
 
+      That->ASTMemUsage = NewAST ? NewAST->getUsedBytes() : 0;
       That->ASTPromise.set_value(
           std::make_shared<ParsedASTWrapper>(std::move(NewAST)));
     } // unlock Mutex
 
     return Diagnostics;
   };
 
   return BindWithForward(FinishRebuild, std::move(Inputs));
std::shared_ptr<const PreambleData> CppFile::getPossiblyStalePreamble() const {
   return LatestAvailablePreamble;
 }
 
 std::shared_future<std::shared_ptr<ParsedASTWrapper>> CppFile::getAST() const {
   std::lock_guard<std::mutex> Lock(Mutex);
   return ASTFuture;
 }
 
+std::size_t CppFile::getUsedBytes() const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  // FIXME: We should not store extra size fields. When we store AST and
+  // Preamble directly, not inside futures, we could compute the sizes from the
+  // stored AST and the preamble in this function directly.
+  return ASTMemUsage + PreambleMemUsage;
+}
+
 CppFile::RebuildGuard::RebuildGuard(CppFile &File,
                                     unsigned RequestRebuildCounter)
     : File(File), RequestRebuildCounter(RequestRebuildCounter) {
   std::unique_lock<std::mutex> Lock(File.Mutex);
   WasCancelledBeforeConstruction = File.RebuildCounter != RequestRebuildCounter;
   if (WasCancelledBeforeConstruction)
     return;
 

clangd/ClangdUnitStore.h

if (It == OpenedFiles.end()) {
       It = OpenedFiles
                .try_emplace(File, CppFile::Create(File, StorePreamblesInMemory,
                                                   std::move(PCHs), ASTCallback))
                .first;
     }
     return It->second;
   }
 
-  std::shared_ptr<CppFile> getFile(PathRef File) {
+  std::shared_ptr<CppFile> getFile(PathRef File) const {
     std::lock_guard<std::mutex> Lock(Mutex);
 
     auto It = OpenedFiles.find(File);
     if (It == OpenedFiles.end())
       return nullptr;
     return It->second;
   }
 
   /// Removes a CppFile, stored for \p File, if it's inside collection and
   /// returns it.
   std::shared_ptr<CppFile> removeIfPresent(PathRef File);
 
+  /// Gets used memory for each of the stored files.
+  std::vector<std::pair<Path, std::size_t>> getUsedBytesPerFile() const;
+
 private:
-  std::mutex Mutex;
+  mutable std::mutex Mutex;
   llvm::StringMap<std::shared_ptr<CppFile>> OpenedFiles;
   ASTParsedCallback ASTCallback;
 };
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/ClangdUnitStore.cpp

std::shared_ptr<CppFile> CppFileCollection::removeIfPresent(PathRef File) {
   auto It = OpenedFiles.find(File);
   if (It == OpenedFiles.end())
     return nullptr;
 
   std::shared_ptr<CppFile> Result = It->second;
   OpenedFiles.erase(It);
   return Result;
 }
+std::vector<std::pair<Path, std::size_t>>
+CppFileCollection::getUsedBytesPerFile() const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  std::vector<std::pair<Path, std::size_t>> Result;
+  Result.reserve(OpenedFiles.size());
+  for (auto &&PathAndFile : OpenedFiles)
+    Result.push_back(
+        {PathAndFile.first().str(), PathAndFile.second->getUsedBytes()});
+  return Result;
+}

unittests/clangd/ClangdTests.cpp

 #include "Context.h"
 #include "TestFS.h"
 #include "clang/Config/config.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Regex.h"
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <algorithm>
 #include <chrono>
 #include <iostream>
 #include <random>
 #include <string>
 #include <thread>
 #include <vector>
 
 namespace clang {
 namespace clangd {
+
+using ::testing::ElementsAre;
+using ::testing::Eq;
+using ::testing::Gt;
+using ::testing::IsEmpty;
+using ::testing::Pair;
+using ::testing::UnorderedElementsAre;
+
 namespace {
 
 // Don't wait for async ops in clangd test more than that to avoid blocking
 // indefinitely in case of bugs.
 static const std::chrono::seconds DefaultFutureTimeout =
     std::chrono::seconds(10);
 
 static bool diagsContainErrors(ArrayRef<DiagWithFixIts> Diagnostics) {
)cpp";
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
   // Subsequent addDocument calls should finish without errors too.
   Server.addDocument(Context::empty(), FooCpp, SourceContents1);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
   Server.addDocument(Context::empty(), FooCpp, SourceContents2);
   EXPECT_FALSE(DiagConsumer.hadErrorInLastDiags());
 }
 
+TEST_F(ClangdVFSTest, MemoryUsage) {
+  MockFSProvider FS;
+  ErrorCheckingDiagConsumer DiagConsumer;
+  MockCompilationDatabase CDB;
+  ClangdServer Server(CDB, DiagConsumer, FS,
+                      /*AsyncThreadsCount=*/0,
+                      /*StorePreamblesInMemory=*/true);
+
+  // No need to sync reparses, because reparses are performed on the calling
+  // thread.
+  Path FooCpp = getVirtualTestFilePath("foo.cpp").str();
+  const auto SourceContents = R"cpp(
+struct Something {
+  int method();
+};
+)cpp";
+  Path BarCpp = getVirtualTestFilePath("bar.cpp").str();
+
+  FS.Files[FooCpp] = "";
+  FS.Files[BarCpp] = "";
+
+  EXPECT_THAT(Server.getUsedBytesPerFile(), IsEmpty());
+
+  Server.addDocument(Context::empty(), FooCpp, SourceContents);
+  Server.addDocument(Context::empty(), BarCpp, SourceContents);
+
+  EXPECT_THAT(Server.getUsedBytesPerFile(),
+              UnorderedElementsAre(Pair(FooCpp, Gt(0u)), Pair(BarCpp, Gt(0u))));
+
+  Server.removeDocument(Context::empty(), FooCpp);
+  EXPECT_THAT(Server.getUsedBytesPerFile(), ElementsAre(Pair(BarCpp, Gt(0u))));
+
+  Server.removeDocument(Context::empty(), BarCpp);
+  EXPECT_THAT(Server.getUsedBytesPerFile(), IsEmpty());
+}
+
 class ClangdThreadingTest : public ClangdVFSTest {};
 
 TEST_F(ClangdThreadingTest, StressTest) {
   // Without 'static' clang gives an error for a usage inside TestDiagConsumer.
   static const unsigned FilesCount = 5;
   const unsigned RequestsCount = 500;
   // Blocking requests wait for the parsing to complete, they slow down the test
   // dramatically, so they are issued rarely. Each