[clangd] Keep only a limited number of idle ASTs in memory

[ilya-biryukov]

Summary

After this commit, clangd will only keep the last 3 accessed ASTs in
memory. Preambles for each of the opened files are still kept in
memory to make completion and AST rebuilds fast.

AST rebuilds are usually fast enough, but having the last ASTs in
memory still considerably improves latency of operations like
findDefinition and documeneHighlight, which are often sent multiple
times a second when moving around the code. So keeping some of the last
accessed ASTs in memory seems like a reasonable tradeoff.

[ilya-biryukov]

Another alternative that I've considered was evicting the ASTs from memory after a certain period of time, e.g. after 30 seconds of inactivity for some file. This might be simpler and also cover the use-case of speeding up multiple code navigation requests (findDefinition/documentHighlight) in a row.
Yet another thing that came to mind: walking over all of the AST to find reference under the cursor is terribly inefficient, we should be able to build a cheap data structure that speeds up this operation. Maybe we can even store enough information to not need the AST for navigation at all and build it only for features like refactorings.
@sammccall, let me know what are your thoughts on all of this.

[malaperle]

In D47063#1104417, @ilya-biryukov wrote:

Maybe we can even store enough information to not need the AST for navigation at all and build it only for features like refactorings.
@sammccall, let me know what are your thoughts on all of this.

That's what I was thinking. I think I also had similar discussion with @arphaman. I think storing a limited number of ASTs is a good interim solution.

[sammccall]

(Haven't reviewed the code yet, just design stuff)
I'm tempted to push back on the idea that we need to store any - "if it's fast enough for code complete, it's fast enough for GTD". But that's probably oversimplifying - we don't force reuse of a stable preamble for GTD I think. So we probably do want some cache.

In D47063#1104495, @malaperle wrote:

In D47063#1104417, @ilya-biryukov wrote:

Maybe we can even store enough information to not need the AST for navigation at all and build it only for features like refactorings.
@sammccall, let me know what are your thoughts on all of this.

That's what I was thinking. I think I also had similar discussion with @arphaman. I think storing a limited number of ASTs is a good interim solution.

Agree - it's a good idea but we shouldn't block on it. A general-purpose xrefs index may solve this problem (and others) but requires a bunch of design. A narrower structure risks building a bunch of complexity for one feature we can't reuse for the next.

Another alternative that I've considered was evicting the ASTs from memory after a certain period of time, e.g. after 30 seconds of inactivity for some file.

We discussed this a bit more offline. A time-based limit reduces idle RAM usage, and a weight limit (or just count) reduces peak RAM.
Relatively speaking, idle seems to be more important to our hosted/multiplexed use cases, and peak is more important when running on a developer machine.
Weight is probably easier to tune. Time is easier to implement as the TUs are independent.

So this gives us some options (assuming we want some caching, but not unlimited):

• Evict if old - this helps hosted a lot, and is simple to implement.
• Evict if full (this patch) - this helps standalone, and is more complex.
• Evict if full AND old - this can be tuned nicely for hosted and standalone. Most complex, but only a little more than the previous option.
• Evict if full OR old - this puts a hard limit on resource usage and controls idle, but doesn't seem useful for hosted as it can't drive idle to zero.

I think the main design decision is whether the cache logic requires TUs to interact (vs simple time eviction). If we pay that complexity cost we get a lot of flexibility to tweak the cache. It's the hosted stuff that's driving this (for Google) right now, but maybe that's just because we're not really measuring impact on workstations yet :)
So I think I like this policy as a starting point, but we should plan to bolt on time-based-expiry in the near future.

[sammccall]

Having taken a closer look, I think the cache can be simplified/separated a bit more cleanly by returning shared pointers and not allowing lookups, instead restoring limited ownership in CppFile...

Happy to discuss more, esp if you might disagree :)

clangd/ClangdUnit.h

132 ↗

(On Diff #147498)

This may be change aversion, but I'm not sure this class does enough after this change - it doesn't store the inputs or the outputs/cache, so it kind of seems like it wants to be a function.

I guess the motivation here is that storing the outputs means dealing with the cache, and the cache is local to TUScheduler.
But CppFile is only used in TUScheduler, so we could move this there too? It feels like expanding the scope more than I'd like.

The upside is that I think it's a more natural division of responsibility: CppFile could continue to be the "main" holder of the shared_ptr<Preamble> (which we don't limit, but share), and instead of Optional<ParsedAST> it'd have a weak_ptr<ParsedAST> which is controlled and can be refreshed through the cache.

As discussed offline, the cache could look something like:

class Cache {
   shared_ptr<ParsedAST> put(ParsedAST);
   void hintUsed(ParsedAST*); // optional, bumps LRU when client reads
   void hintUnused(ParsedAST*); // optional, releases when client abandons
}

shared_ptr<ParsedAST> CppFile::getAST() {
  shared_ptr<ParsedAST> AST = WeakAST.lock();
  if (AST)
    Cache.hintUsed(AST.get());
  else
    WeakAST = AST = Cache.put(build...);
  return AST;
}

[ilya-biryukov]

• Reimplemented LRU cache with shared_ptr and weak_ptr.

[ilya-biryukov]

clangd/ClangdUnit.h

132 ↗

(On Diff #147498)

I've reimplemented the cache with weak_ptr/shared_ptr.
With a slightly different interface to hide more stuff in the cache API. Please take a look and let me know what you think.

Nevertheless, I still find CppFile refactoring useful. It unties preambles from the ASTs and I believe this is a good thing, given that their lifetimes are different.

[sammccall]

Thanks, this looks a lot better!
There do seem to be a lot of little classes that exist exactly one-per-TU (ASTWorker, ASTBuilder, CachedAST, to a lesser extent ParsedAST) and I'm not sure the balance of responsibilities is quite right. Some comments below.

clangd/ClangdUnit.h

132 ↗

(On Diff #147498)

I'm not sure how much they were tangled before, they were computed in the same place, but accessed separately, and it's not sure you ever *want* to compute them separately? (possibly in unit tests?)

I do think making ASTWorker maintain the old preamble and pass it in is confusing. The remaining members are trivial and unrelated enough that I do think if the references to the preamble/ast are removed, then moving the remaining members to ASTWorker or to a dumb struct, and making these free functions would make it easier to navigate the class structure.

clangd/TUScheduler.cpp

71 ↗

(On Diff #148355)

I think I understand this more as "updates the value" but the value is lazy...

I find this API somewhat hard to follow, maybe just because it's unfamiliar.
I've mostly seen cache APIs look like one of:

1. Cache(function<Value(Input)> Compute), Value Cache::get(Input)
2. void Cache::put(Key, Value), Optional<Value> Cache::get(Key)
3. Handle Cache::put(Value), Optional<Value> Handle::get()

This one is Slot Cache::create(), void Slot::update(function<Value()> LazyV), Value Slot::get().

It's similar-ish to 3, but has 3 nontrivial operations instead of 2, and the slot object is externally mutable instead of immutable, so it seems more complex. What does it buy us in exchange?

(1 & 2 work well with a natural key or inputs that are easy to compare, which we don't particularly have)

91 ↗

(On Diff #148355)

naming: IdleASTs doesn't mention the function of this class, which is a cache.
I'd suggest swapping the names: call the class ASTCache and the *instance* IdleASTs as that reflects its role within the TUScheduler.

For CachedAST, I think the relationship would be well-exposed by nesting it as ASTCache::Entry. This also gives you the friend for free, which seems like a hint that it's an appropriate structure.
(Though I'm not sure CachedAST is that useful)

153 ↗

(On Diff #148355)

as discussed offline, using a CachedAST* as a key shouldn't be necessary, the ParsedAST* should be enough I think.

157 ↗

(On Diff #148355)

document why

clangd/TUScheduler.h

48 ↗

(On Diff #148355)

nit: Policy would be more specific than Params

[ilya-biryukov]

• Rebase, fix merge conflicts
• Simpler implemenataion of the Cache
• s/IdleASTs/ASTCache

[ilya-biryukov]

I have addressed the comments regarding the cache implementation. ASTBuilder ones are still pending, but I would appreciate the feedback on how TUScheduler.cpp looks like.

clangd/ClangdUnit.h

132 ↗

(On Diff #147498)

it's not sure you ever *want* to compute them separately?

I think that's exactly what we're doing now. The ASTs can now get built separately from the preamble, because they can be evicted from the cache and need to be rebuild while preamble is not changed.

The remaining members are trivial and unrelated enough that I do think if the references to the preamble/ast are removed, then moving the remaining members to ASTWorker or to a dumb struct, and making these free functions would make it easier to navigate the class structure.

Dumb struct SG, it's essentially what ASTWorker is right now.

clangd/TUScheduler.cpp

71 ↗

(On Diff #148355)

As discussed offline, now we have a simpler version that keeps unique_ptrs to idle ASTs and the clients are responsible for building the ASTs.
Note that it's not a "cache" per se, so we might want a different name for that.
@sammccall, you suggested to call it a pool, I find it reasonable. Should we name it ASTPool instead of ASTCache?

[sammccall]

The cache looks way simpler now, thank you!
As discussed offline, flattening ASTBuilder right into ASTWorker still seems like a good idea to me, but happy with what you come up with there.

clangd/TUScheduler.cpp

66 ↗

(On Diff #148599)

I'd say a little more about the interaction here. e.g.

/// An LRU cache of idle ASTs.
/// Because we want to limit the overall number of these we retain, the cache
/// owns ASTs (and may evict them) while their workers are idle.
/// Workers borrow them when active, and return them when done.

84 ↗

(On Diff #148599)

consider assert(findByKey(K) == LRU.end()) as a precondition

92 ↗

(On Diff #148599)

Just "run the expensive destructor outside the lock"?
the "may not" case seems unimportant and slightly confusing here

94 ↗

(On Diff #148599)

this line isn't actually needed right?

342 ↗

(On Diff #148599)

This failure doesn't get cached, correct? That's bad for performance.

But if we think this is always a clangd bug, it's probably fine. (and certainly simplifies things)

71 ↗

(On Diff #148355)

I think the name is actually fine, it's still mostly a cache.
It does have things in common with a pool, but unrelated consumers can't share a resource, so I think that name is at least as misleading.

[ilya-biryukov]

• Address review comments.

[ilya-biryukov]

• Remove ASTBuilder, make everything a helper function

[ilya-biryukov]

clangd/TUScheduler.cpp

94 ↗

(On Diff #148599)

It isn't. But it makes an important operation (destructor of the AST) explicit, so I'd still keep it.

342 ↗

(On Diff #148599)

Thanks, good catch! I somehow missed it, since at some point the failure was cached.
I think we should cache it.

Failing ASTs is a clangd bug, of course.
However, they might be hard to fix if it's something inside clang, so I believe we should handle the failures gracefully in that case.

71 ↗

(On Diff #148355)

SG, let's leave as is.

[ilya-biryukov]

• s/Params/Policy

[ilya-biryukov]

There are still a few nits I haven't addressed, but the other big change is now there: removing ASTBuilder, replacing it with free-standing functions. I'd say I liked the previous code better, since the use sites were a bit smaller and the functions have ridiculously many params now. But all of that seems manageable at this point.
@sammccall, let me know what you think. And I'll address the nits that are left tomorrow.

[sammccall]

Impl looks good.
Is there a way we can reasonably test this? (specifically that we don't just store all the ASTs forever)

clangd/ClangdUnit.h

143 ↗

(On Diff #148824)

nit: i think filename here is only used for logging, just use Inputs.CompileCommand.Filename?

clangd/TUScheduler.cpp

103 ↗

(On Diff #148824)

please also explicitly mention what nullptr means (even though that's really up to the caller of put()). None vs nullptr leaves room for confusion.

clangd/TUScheduler.h

66 ↗

(On Diff #148824)

does this actually have more than one caller? what's the plan for exposing this option to embedders/CLI users (not saying we necessarily need the latter)?

[ilya-biryukov]

• Added a unit test
• Address review comments
• Add ASTRetentionPolicy param to ClangdServer

clangd/ClangdUnit.h

143 ↗

(On Diff #148824)

Tried doing that, but the filename parameter is actually passed to PreambleCallback that updates the dynamic index. Using filename from compile command there seems fragile, so I kept the parameter for now.

clangd/TUScheduler.h

66 ↗

(On Diff #148824)

Yes, just one caller outside the tests.
The plan was to expose it only in ClangdServer for now. Giving this knob in CLI might be useful, if we have good reasons for that, but I hope that we could pick the default that work for everyone instead.
Added that as a parameter of ClangdServer.

Maybe we should move the default value of 3 to ClangdServer? WDYT?

[ilya-biryukov]

• Fixed formatting

[sammccall]

clangd/TUScheduler.h

66 ↗

(On Diff #148824)

CLI when needed SG.
I think we want to give our cloud folks the chance to set it to zero.
So maybe put the default in ClangdLSPServer? (or make it a param there too and move the value to main)

[ilya-biryukov]

clangd/TUScheduler.h

66 ↗

(On Diff #148824)

After an online chat we've decided to keep the default inside the ASTRetentionPolicy struct.
If we add the CLI parameter, we can pass custom values from main to ClangdLSPServer.
And other clients already have a knob to set it in ClangdServer.

clang-tools-extra/trunk/clangd/ClangdServer.h

 /// definition.
 class ClangdServer {
 public:
   struct Options {
     /// To process requests asynchronously, ClangdServer spawns worker threads.
     /// If 0, all requests are processed on the calling thread.
     unsigned AsyncThreadsCount = getDefaultAsyncThreadsCount();
 
+    /// AST caching policy. The default is to keep up to 3 ASTs in memory.
+    ASTRetentionPolicy RetentionPolicy;
+
     /// Cached preambles are potentially large. If false, store them on disk.
     bool StorePreamblesInMemory = true;
 
     /// If true, ClangdServer builds a dynamic in-memory index for symbols in
     /// opened files and uses the index to augment code completion results.
     bool BuildDynamicSymbolIndex = false;
 
     /// If set, use this index to augment code completion results.

clang-tools-extra/trunk/clangd/ClangdServer.cpp

: CompileArgs(CDB, Opts.ResourceDir ? Opts.ResourceDir->str()
       // critical paths.
       WorkScheduler(
           Opts.AsyncThreadsCount, Opts.StorePreamblesInMemory,
           FileIdx
               ? [this](PathRef Path, ASTContext &AST,
                        std::shared_ptr<Preprocessor>
                            PP) { FileIdx->update(Path, &AST, std::move(PP)); }
               : PreambleParsedCallback(),
-          Opts.UpdateDebounce) {
+          Opts.UpdateDebounce, Opts.RetentionPolicy) {
   if (FileIdx && Opts.StaticIndex) {
     MergedIndex = mergeIndex(FileIdx.get(), Opts.StaticIndex);
     Index = MergedIndex.get();
   } else if (FileIdx)
     Index = FileIdx.get();
   else if (Opts.StaticIndex)
     Index = Opts.StaticIndex;
   else

clang-tools-extra/trunk/clangd/ClangdUnit.h

 
 namespace clangd {
 
 // Stores Preamble and associated data.
 struct PreambleData {
   PreambleData(PrecompiledPreamble Preamble, std::vector<Diag> Diags,
                std::vector<Inclusion> Inclusions);
 
+  tooling::CompileCommand CompileCommand;
   PrecompiledPreamble Preamble;
   std::vector<Diag> Diags;
   // Processes like code completions and go-to-definitions will need #include
   // information, and their compile action skips preamble range.
   std::vector<Inclusion> Inclusions;
 };
 
 /// Information required to run clang, e.g. to parse AST or do code completion.
private:
   // top-level decls from the preamble.
   std::vector<const Decl *> LocalTopLevelDecls;
   std::vector<Inclusion> Inclusions;
 };
 
 using PreambleParsedCallback = std::function<void(
     PathRef Path, ASTContext &, std::shared_ptr<clang::Preprocessor>)>;
 
-/// Manages resources, required by clangd. Allows to rebuild file with new
-/// contents, and provides AST and Preamble for it.
-class CppFile {
-public:
-  CppFile(PathRef FileName, bool StorePreamblesInMemory,
-          std::shared_ptr<PCHContainerOperations> PCHs,
+/// Builds compiler invocation that could be used to build AST or preamble.
+std::unique_ptr<CompilerInvocation>
+buildCompilerInvocation(const ParseInputs &Inputs);
+
+/// Rebuild the preamble for the new inputs unless the old one can be reused.
+/// If \p OldPreamble can be reused, it is returned unchanged.
+/// If \p OldPreamble is null, always builds the preamble.
+/// If \p PreambleCallback is set, it will be run on top of the AST while
+/// building the preamble. Note that if the old preamble was reused, no AST is
+/// built and, therefore, the callback will not be executed.
+std::shared_ptr<const PreambleData>
+buildPreamble(PathRef FileName, CompilerInvocation &CI,
+              std::shared_ptr<const PreambleData> OldPreamble,
+              const tooling::CompileCommand &OldCompileCommand,
+              const ParseInputs &Inputs,
+              std::shared_ptr<PCHContainerOperations> PCHs, bool StoreInMemory,
               PreambleParsedCallback PreambleCallback);
 
-  /// Rebuild the AST and the preamble.
-  /// Returns a list of diagnostics or llvm::None, if an error occured.
-  llvm::Optional<std::vector<Diag>> rebuild(ParseInputs &&Inputs);
-  /// Returns the last built preamble.
-  const std::shared_ptr<const PreambleData> &getPreamble() const;
-  /// Returns the last built AST.
-  ParsedAST *getAST() const;
-  /// Returns an estimated size, in bytes, currently occupied by the AST and the
-  /// Preamble.
-  std::size_t getUsedBytes() const;
-
-private:
-  /// Build a new preamble for \p Inputs. If the current preamble can be reused,
-  /// it is returned instead.
-  /// This method is const to ensure we don't incidentally modify any fields.
-  std::shared_ptr<const PreambleData>
-  rebuildPreamble(CompilerInvocation &CI,
-                  const tooling::CompileCommand &Command,
-                  IntrusiveRefCntPtr<vfs::FileSystem> FS,
-                  llvm::MemoryBuffer &ContentsBuffer) const;
-
-  const Path FileName;
-  const bool StorePreamblesInMemory;
-
-  /// The last CompileCommand used to build AST and Preamble.
-  tooling::CompileCommand Command;
-  /// The last parsed AST.
-  llvm::Optional<ParsedAST> AST;
-  /// The last built Preamble.
-  std::shared_ptr<const PreambleData> Preamble;
-  /// Utility class required by clang
-  std::shared_ptr<PCHContainerOperations> PCHs;
-  /// This is called after the file is parsed. This can be nullptr if there is
-  /// no callback.
-  PreambleParsedCallback PreambleCallback;
-};
+/// Build an AST from provided user inputs. This function does not check if
+/// preamble can be reused, as this function expects that \p Preamble is the
+/// result of calling buildPreamble.
+llvm::Optional<ParsedAST>
+buildAST(PathRef FileName, std::unique_ptr<CompilerInvocation> Invocation,
+         const ParseInputs &Inputs,
+         std::shared_ptr<const PreambleData> Preamble,
+         std::shared_ptr<PCHContainerOperations> PCHs);
 
 /// Get the beginning SourceLocation at a specified \p Pos.
 /// May be invalid if Pos is, or if there's no identifier.
 SourceLocation getBeginningOfIdentifier(ParsedAST &Unit, const Position &Pos,
                                         const FileID FID);
 
 /// For testing/debugging purposes. Note that this method deserializes all
 /// unserialized Decls, so use with care.
 void dumpAST(ParsedAST &AST, llvm::raw_ostream &OS);
 
 } // namespace clangd
 } // namespace clang
 #endif

clang-tools-extra/trunk/clangd/ClangdUnit.cpp

ParsedAST::Build(std::unique_ptr<clang::CompilerInvocation> CI,
 
   // UnitDiagsConsumer is local, we can not store it in CompilerInstance that
   // has a longer lifetime.
   Clang->getDiagnostics().setClient(new IgnoreDiagnostics);
   // CompilerInstance won't run this callback, do it directly.
   ASTDiags.EndSourceFile();
 
   std::vector<const Decl *> ParsedDecls = Action->takeTopLevelDecls();
+  std::vector<Diag> Diags = ASTDiags.take();
+  // Add diagnostics from the preamble, if any.
+  if (Preamble)
+    Diags.insert(Diags.begin(), Preamble->Diags.begin(), Preamble->Diags.end());
   return ParsedAST(std::move(Preamble), std::move(Clang), std::move(Action),
-                   std::move(ParsedDecls), ASTDiags.take(),
+                   std::move(ParsedDecls), std::move(Diags),
                    std::move(Inclusions));
 }
 
 ParsedAST::ParsedAST(ParsedAST &&Other) = default;
 
 ParsedAST &ParsedAST::operator=(ParsedAST &&Other) = default;
 
 ParsedAST::~ParsedAST() {
ParsedAST::ParsedAST(std::shared_ptr<const PreambleData> Preamble,
     : Preamble(std::move(Preamble)), Clang(std::move(Clang)),
       Action(std::move(Action)), Diags(std::move(Diags)),
       LocalTopLevelDecls(std::move(LocalTopLevelDecls)),
       Inclusions(std::move(Inclusions)) {
   assert(this->Clang);
   assert(this->Action);
 }
 
-CppFile::CppFile(PathRef FileName, bool StorePreamblesInMemory,
-                 std::shared_ptr<PCHContainerOperations> PCHs,
-                 PreambleParsedCallback PreambleCallback)
-    : FileName(FileName), StorePreamblesInMemory(StorePreamblesInMemory),
-      PCHs(std::move(PCHs)), PreambleCallback(std::move(PreambleCallback)) {
-  log("Created CppFile for " + FileName);
-}
-
-llvm::Optional<std::vector<Diag>> CppFile::rebuild(ParseInputs &&Inputs) {
-  log("Rebuilding file " + FileName + " with command [" +
-      Inputs.CompileCommand.Directory + "] " +
-      llvm::join(Inputs.CompileCommand.CommandLine, " "));
-
+std::unique_ptr<CompilerInvocation>
+clangd::buildCompilerInvocation(const ParseInputs &Inputs) {
   std::vector<const char *> ArgStrs;
   for (const auto &S : Inputs.CompileCommand.CommandLine)
     ArgStrs.push_back(S.c_str());
 
   if (Inputs.FS->setCurrentWorkingDirectory(Inputs.CompileCommand.Directory)) {
-    log("Couldn't set working directory");
-    // We run parsing anyway, our lit-tests rely on results for non-existing
-    // working dirs.
+    log("Couldn't set working directory when creating compiler invocation.");
+    // We proceed anyway, our lit-tests rely on results for non-existing working
+    // dirs.
   }
 
-  // Prepare CompilerInvocation.
-  std::unique_ptr<CompilerInvocation> CI;
-  {
-    // FIXME(ibiryukov): store diagnostics from CommandLine when we start
-    // reporting them.
-    IgnoreDiagnostics IgnoreDiagnostics;
-    IntrusiveRefCntPtr<DiagnosticsEngine> CommandLineDiagsEngine =
-        CompilerInstance::createDiagnostics(new DiagnosticOptions,
-                                            &IgnoreDiagnostics, false);
-    CI = createInvocationFromCommandLine(ArgStrs, CommandLineDiagsEngine,
-                                         Inputs.FS);
-    if (!CI) {
-      log("Could not build CompilerInvocation for file " + FileName);
-      AST = llvm::None;
-      Preamble = nullptr;
-      return llvm::None;
-    }
-    // createInvocationFromCommandLine sets DisableFree.
-    CI->getFrontendOpts().DisableFree = false;
-    CI->getLangOpts()->CommentOpts.ParseAllComments = true;
-  }
+  std::unique_ptr<CompilerInvocation> CI = createInvocationFromCommandLine(
+      ArgStrs, CommandLineDiagsEngine, Inputs.FS);
+  if (!CI)
+    return nullptr;
+  return CI;
 
-  std::unique_ptr<llvm::MemoryBuffer> ContentsBuffer =
-      llvm::MemoryBuffer::getMemBufferCopy(Inputs.Contents, FileName);
-
-  // Compute updated Preamble.
-  std::shared_ptr<const PreambleData> NewPreamble =
-      rebuildPreamble(*CI, Inputs.CompileCommand, Inputs.FS, *ContentsBuffer);
-
-  // Remove current AST to avoid wasting memory.
-  AST = llvm::None;
-  // Compute updated AST.
-  llvm::Optional<ParsedAST> NewAST;
-  {
-    trace::Span Tracer("Build");
-    SPAN_ATTACH(Tracer, "File", FileName);
-    NewAST = ParsedAST::Build(std::move(CI), NewPreamble,
-                              std::move(ContentsBuffer), PCHs, Inputs.FS);
-  }
-
-  std::vector<Diag> Diagnostics;
-  if (NewAST) {
-    // Collect diagnostics from both the preamble and the AST.
-    if (NewPreamble)
-      Diagnostics = NewPreamble->Diags;
-    Diagnostics.insert(Diagnostics.end(), NewAST->getDiagnostics().begin(),
-                       NewAST->getDiagnostics().end());
-  }
-
-  // Write the results of rebuild into class fields.
-  Command = std::move(Inputs.CompileCommand);
-  Preamble = std::move(NewPreamble);
-  AST = std::move(NewAST);
-  return Diagnostics;
-}
-
-const std::shared_ptr<const PreambleData> &CppFile::getPreamble() const {
-  return Preamble;
-}
-
-ParsedAST *CppFile::getAST() const {
-  // We could add mutable to AST instead of const_cast here, but that would also
-  // allow writing to AST from const methods.
-  return AST ? const_cast<ParsedAST *>(AST.getPointer()) : nullptr;
-}
-
-std::size_t CppFile::getUsedBytes() const {
-  std::size_t Total = 0;
-  if (AST)
-    Total += AST->getUsedBytes();
-  if (StorePreamblesInMemory && Preamble)
-    Total += Preamble->Preamble.getSize();
-  return Total;
-}
-
-std::shared_ptr<const PreambleData>
-CppFile::rebuildPreamble(CompilerInvocation &CI,
-                         const tooling::CompileCommand &Command,
-                         IntrusiveRefCntPtr<vfs::FileSystem> FS,
-                         llvm::MemoryBuffer &ContentsBuffer) const {
-  const auto &OldPreamble = this->Preamble;
-  auto Bounds = ComputePreambleBounds(*CI.getLangOpts(), &ContentsBuffer, 0);
-  if (OldPreamble && compileCommandsAreEqual(this->Command, Command) &&
-      OldPreamble->Preamble.CanReuse(CI, &ContentsBuffer, Bounds, FS.get())) {
+std::shared_ptr<const PreambleData> clangd::buildPreamble(
+    PathRef FileName, CompilerInvocation &CI,
+    std::shared_ptr<const PreambleData> OldPreamble,
+    const tooling::CompileCommand &OldCompileCommand, const ParseInputs &Inputs,
+    std::shared_ptr<PCHContainerOperations> PCHs, bool StoreInMemory,
+    PreambleParsedCallback PreambleCallback) {
+  // Note that we don't need to copy the input contents, preamble can live
+  // without those.
+  auto ContentsBuffer = llvm::MemoryBuffer::getMemBuffer(Inputs.Contents);
+  auto Bounds =
+      ComputePreambleBounds(*CI.getLangOpts(), ContentsBuffer.get(), 0);
+
+  if (OldPreamble &&
+      compileCommandsAreEqual(Inputs.CompileCommand, OldCompileCommand) &&
+      OldPreamble->Preamble.CanReuse(CI, ContentsBuffer.get(), Bounds,
+                                     Inputs.FS.get())) {
     log("Reusing preamble for file " + Twine(FileName));
     return OldPreamble;
   }
   log("Preamble for file " + Twine(FileName) +
       " cannot be reused. Attempting to rebuild it.");
 
-  trace::Span Tracer("Preamble");
+  trace::Span Tracer("BuildPreamble");
   SPAN_ATTACH(Tracer, "File", FileName);
   StoreDiags PreambleDiagnostics;
   IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
       CompilerInstance::createDiagnostics(&CI.getDiagnosticOpts(),
                                           &PreambleDiagnostics, false);
 
   // 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;
 
   CppFilePreambleCallbacks SerializedDeclsCollector(FileName, PreambleCallback);
+  if (Inputs.FS->setCurrentWorkingDirectory(Inputs.CompileCommand.Directory)) {
+    log("Couldn't set working directory when building the preamble.");
+    // We proceed anyway, our lit-tests rely on results for non-existing working
+    // dirs.
+  }
   auto BuiltPreamble = PrecompiledPreamble::Build(
-      CI, &ContentsBuffer, Bounds, *PreambleDiagsEngine, FS, PCHs,
-      /*StoreInMemory=*/StorePreamblesInMemory, SerializedDeclsCollector);
+      CI, ContentsBuffer.get(), Bounds, *PreambleDiagsEngine, Inputs.FS, PCHs,
+      StoreInMemory, SerializedDeclsCollector);
 
   // When building the AST for the main file, we do want the function
   // bodies.
   CI.getFrontendOpts().SkipFunctionBodies = false;
 
   if (BuiltPreamble) {
     log("Built preamble of size " + Twine(BuiltPreamble->getSize()) +
         " for file " + Twine(FileName));
-
     return std::make_shared<PreambleData>(
         std::move(*BuiltPreamble), PreambleDiagnostics.take(),
         SerializedDeclsCollector.takeInclusions());
   } else {
     log("Could not build a preamble for file " + Twine(FileName));
     return nullptr;
   }
 }
 
+llvm::Optional<ParsedAST> clangd::buildAST(
+    PathRef FileName, std::unique_ptr<CompilerInvocation> Invocation,
+    const ParseInputs &Inputs, std::shared_ptr<const PreambleData> Preamble,
+    std::shared_ptr<PCHContainerOperations> PCHs) {
+  trace::Span Tracer("BuildAST");
+  SPAN_ATTACH(Tracer, "File", FileName);
+
+  if (Inputs.FS->setCurrentWorkingDirectory(Inputs.CompileCommand.Directory)) {
+    log("Couldn't set working directory when building the preamble.");
+    // We proceed anyway, our lit-tests rely on results for non-existing working
+    // dirs.
+  }
+
+  return ParsedAST::Build(
+      llvm::make_unique<CompilerInvocation>(*Invocation), Preamble,
+      llvm::MemoryBuffer::getMemBufferCopy(Inputs.Contents), PCHs, Inputs.FS);
+}
+
 SourceLocation clangd::getBeginningOfIdentifier(ParsedAST &Unit,
                                                 const Position &Pos,
                                                 const FileID FID) {
   const ASTContext &AST = Unit.getASTContext();
   const SourceManager &SourceMgr = AST.getSourceManager();
   auto Offset = positionToOffset(SourceMgr.getBufferData(FID), Pos);
   if (!Offset) {
     log("getBeginningOfIdentifier: " + toString(Offset.takeError()));

clang-tools-extra/trunk/clangd/TUScheduler.h

 /// Determines whether diagnostics should be generated for a file snapshot.
 enum class WantDiagnostics {
   Yes,  /// Diagnostics must be generated for this snapshot.
   No,   /// Diagnostics must not be generated for this snapshot.
   Auto, /// Diagnostics must be generated for this snapshot or a subsequent one,
         /// within a bounded amount of time.
 };
 
+/// Configuration of the AST retention policy. This only covers retention of
+/// *idle* ASTs. If queue has operations requiring the AST, they might be
+/// kept in memory.
+struct ASTRetentionPolicy {
+  /// Maximum number of ASTs to be retained in memory when there are no pending
+  /// requests for them.
+  unsigned MaxRetainedASTs = 3;
+};
+
 /// Handles running tasks for ClangdServer and managing the resources (e.g.,
 /// preambles and ASTs) for opened files.
 /// TUScheduler is not thread-safe, only one thread should be providing updates
 /// and scheduling tasks.
 /// Callbacks are run on a threadpool and it's appropriate to do slow work in
 /// them. Each task has a name, used for tracing (should be UpperCamelCase).
 /// FIXME(sammccall): pull out a scheduler options struct.
 class TUScheduler {
 public:
   TUScheduler(unsigned AsyncThreadsCount, bool StorePreamblesInMemory,
               PreambleParsedCallback PreambleCallback,
-              std::chrono::steady_clock::duration UpdateDebounce);
+              std::chrono::steady_clock::duration UpdateDebounce,
+              ASTRetentionPolicy RetentionPolicy);
   ~TUScheduler();
 
   /// Returns estimated memory usage for each of the currently open files.
   /// The order of results is unspecified.
   std::vector<std::pair<Path, std::size_t>> getUsedBytesPerFile() const;
 
+  /// Returns a list of files with ASTs currently stored in memory. This method
+  /// is not very reliable and is only used for test. E.g., the results will not
+  /// contain files that currently run something over their AST.
+  std::vector<Path> getFilesWithCachedAST() const;
+
   /// Schedule an update for \p File. Adds \p File to a list of tracked files if
   /// \p File was not part of it before.
   /// FIXME(ibiryukov): remove the callback from this function.
   void update(PathRef File, ParseInputs Inputs, WantDiagnostics WD,
               UniqueFunction<void(std::vector<Diag>)> OnUpdated);
 
   /// Remove \p File from the list of tracked files and schedule removal of its
   /// resources.
public:
   /// Wait until there are no scheduled or running tasks.
   /// Mostly useful for synchronizing tests.
   bool blockUntilIdle(Deadline D) const;
 
 private:
   /// This class stores per-file data in the Files map.
   struct FileData;
 
+public:
+  /// Responsible for retaining and rebuilding idle ASTs. An implementation is
+  /// an LRU cache.
+  class ASTCache;
+
+private:
   const bool StorePreamblesInMemory;
   const std::shared_ptr<PCHContainerOperations> PCHOps;
   const PreambleParsedCallback PreambleCallback;
   Semaphore Barrier;
   llvm::StringMap<std::unique_ptr<FileData>> Files;
+  std::unique_ptr<ASTCache> IdleASTs;
   // None when running tasks synchronously and non-None when running tasks
   // asynchronously.
   llvm::Optional<AsyncTaskRunner> PreambleTasks;
   llvm::Optional<AsyncTaskRunner> WorkerThreads;
   std::chrono::steady_clock::duration UpdateDebounce;
 };
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/trunk/clangd/TUScheduler.cpp

 // - If subsequent updates come before rebuild was started, we replace the
 //   pending inputs and reset the timer.
 // - If any reads of the AST are scheduled, we start building the AST
 //   immediately.
 
 #include "TUScheduler.h"
 #include "Logger.h"
 #include "Trace.h"
+#include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PCHContainerOperations.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Path.h"
+#include <algorithm>
 #include <memory>
 #include <queue>
 #include <thread>
 
 namespace clang {
 namespace clangd {
 using std::chrono::steady_clock;
+
+namespace {
+class ASTWorker;
+}
+
+/// An LRU cache of idle ASTs.
+/// Because we want to limit the overall number of these we retain, the cache
+/// owns ASTs (and may evict them) while their workers are idle.
+/// Workers borrow ASTs when active, and return them when done.
+class TUScheduler::ASTCache {
+public:
+  using Key = const ASTWorker *;
+
+  ASTCache(unsigned MaxRetainedASTs) : MaxRetainedASTs(MaxRetainedASTs) {}
+
+  /// Returns result of getUsedBytes() for the AST cached by \p K.
+  /// If no AST is cached, 0 is returned.
+  bool getUsedBytes(Key K) {
+    std::lock_guard<std::mutex> Lock(Mut);
+    auto It = findByKey(K);
+    if (It == LRU.end() || !It->second)
+      return 0;
+    return It->second->getUsedBytes();
+  }
+
+  /// Store the value in the pool, possibly removing the last used AST.
+  /// The value should not be in the pool when this function is called.
+  void put(Key K, std::unique_ptr<ParsedAST> V) {
+    std::unique_lock<std::mutex> Lock(Mut);
+    assert(findByKey(K) == LRU.end());
+
+    LRU.insert(LRU.begin(), {K, std::move(V)});
+    if (LRU.size() <= MaxRetainedASTs)
+      return;
+    // We're past the limit, remove the last element.
+    std::unique_ptr<ParsedAST> ForCleanup = std::move(LRU.back().second);
+    LRU.pop_back();
+    // Run the expensive destructor outside the lock.
+    Lock.unlock();
+    ForCleanup.reset();
+  }
+
+  /// Returns the cached value for \p K, or llvm::None if the value is not in
+  /// the cache anymore. If nullptr was cached for \p K, this function will
+  /// return a null unique_ptr wrapped into an optional.
+  llvm::Optional<std::unique_ptr<ParsedAST>> take(Key K) {
+    std::unique_lock<std::mutex> Lock(Mut);
+    auto Existing = findByKey(K);
+    if (Existing == LRU.end())
+      return llvm::None;
+    std::unique_ptr<ParsedAST> V = std::move(Existing->second);
+    LRU.erase(Existing);
+    return V;
+  }
+
+private:
+  using KVPair = std::pair<Key, std::unique_ptr<ParsedAST>>;
+
+  std::vector<KVPair>::iterator findByKey(Key K) {
+    return std::find_if(LRU.begin(), LRU.end(),
+                        [K](const KVPair &P) { return P.first == K; });
+  }
+
+  std::mutex Mut;
+  unsigned MaxRetainedASTs;
+  /// Items sorted in LRU order, i.e. first item is the most recently accessed
+  /// one.
+  std::vector<KVPair> LRU; /* GUARDED_BY(Mut) */
+};
+
 namespace {
 class ASTWorkerHandle;
 
 /// Owns one instance of the AST, schedules updates and reads of it.
 /// Also responsible for building and providing access to the preamble.
 /// Each ASTWorker processes the async requests sent to it on a separate
 /// dedicated thread.
 /// The ASTWorker that manages the AST is shared by both the processing thread
 /// and the TUScheduler. The TUScheduler should discard an ASTWorker when
 /// remove() is called, but its thread may be busy and we don't want to block.
 /// So the workers are accessed via an ASTWorkerHandle. Destroying the handle
 /// signals the worker to exit its run loop and gives up shared ownership of the
 /// worker.
 class ASTWorker {
   friend class ASTWorkerHandle;
-  ASTWorker(llvm::StringRef File, Semaphore &Barrier, CppFile AST, bool RunSync,
-            steady_clock::duration UpdateDebounce);
+  ASTWorker(PathRef FileName, TUScheduler::ASTCache &LRUCache,
+            Semaphore &Barrier, bool RunSync,
+            steady_clock::duration UpdateDebounce,
+            std::shared_ptr<PCHContainerOperations> PCHs,
+            bool StorePreamblesInMemory,
+            PreambleParsedCallback PreambleCallback);
 
 public:
   /// Create a new ASTWorker and return a handle to it.
   /// The processing thread is spawned using \p Tasks. However, when \p Tasks
   /// is null, all requests will be processed on the calling thread
   /// synchronously instead. \p Barrier is acquired when processing each
   /// request, it is be used to limit the number of actively running threads.
-  static ASTWorkerHandle Create(llvm::StringRef File, AsyncTaskRunner *Tasks,
-                                Semaphore &Barrier, CppFile AST,
-                                steady_clock::duration UpdateDebounce);
+  static ASTWorkerHandle Create(PathRef FileName,
+                                TUScheduler::ASTCache &IdleASTs,
+                                AsyncTaskRunner *Tasks, Semaphore &Barrier,
+                                steady_clock::duration UpdateDebounce,
+                                std::shared_ptr<PCHContainerOperations> PCHs,
+                                bool StorePreamblesInMemory,
+                                PreambleParsedCallback PreambleCallback);
   ~ASTWorker();
 
   void update(ParseInputs Inputs, WantDiagnostics,
               UniqueFunction<void(std::vector<Diag>)> OnUpdated);
   void runWithAST(llvm::StringRef Name,
                   UniqueFunction<void(llvm::Expected<InputsAndAST>)> Action);
   bool blockUntilIdle(Deadline Timeout) const;
 
   std::shared_ptr<const PreambleData> getPossiblyStalePreamble() const;
   std::size_t getUsedBytes() const;
+  bool isASTCached() const;
 
 private:
   // Must be called exactly once on processing thread. Will return after
   // stop() is called on a separate thread and all pending requests are
   // processed.
   void run();
   /// Signal that run() should finish processing pending requests and exit.
   void stop();
private:
   struct Request {
     UniqueFunction<void()> Action;
     std::string Name;
     steady_clock::time_point AddTime;
     Context Ctx;
     llvm::Optional<WantDiagnostics> UpdateType;
   };
 
-  const std::string File;
+  /// Handles retention of ASTs.
+  TUScheduler::ASTCache &IdleASTs;
   const bool RunSync;
-  // Time to wait after an update to see whether another update obsoletes it.
+  /// Time to wait after an update to see whether another update obsoletes it.
   const steady_clock::duration UpdateDebounce;
+  /// File that ASTWorker is reponsible for.
+  const Path FileName;
+  /// Whether to keep the built preambles in memory or on disk.
+  const bool StorePreambleInMemory;
+  /// Callback, passed to the preamble builder.
+  const PreambleParsedCallback PreambleCallback;
+  /// Helper class required to build the ASTs.
+  const std::shared_ptr<PCHContainerOperations> PCHs;
 
   Semaphore &Barrier;
-  // AST and FileInputs are only accessed on the processing thread from run().
-  CppFile AST;
-  // Inputs, corresponding to the current state of AST.
+  /// Inputs, corresponding to the current state of AST.
   ParseInputs FileInputs;
-  // Guards members used by both TUScheduler and the worker thread.
+  /// CompilerInvocation used for FileInputs.
+  std::unique_ptr<CompilerInvocation> Invocation;
+  /// Size of the last AST
+  /// Guards members used by both TUScheduler and the worker thread.
   mutable std::mutex Mutex;
   std::shared_ptr<const PreambleData> LastBuiltPreamble; /* GUARDED_BY(Mutex) */
-  // Result of getUsedBytes() after the last rebuild or read of AST.
-  std::size_t LastASTSize; /* GUARDED_BY(Mutex) */
-  // Set to true to signal run() to finish processing.
+  /// Set to true to signal run() to finish processing.
   bool Done;                    /* GUARDED_BY(Mutex) */
   std::deque<Request> Requests; /* GUARDED_BY(Mutex) */
   mutable std::condition_variable RequestsCV;
 };
 
 /// A smart-pointer-like class that points to an active ASTWorker.
 /// In destructor, signals to the underlying ASTWorker that no new requests will
 /// be sent and the processing loop may exit (after running all pending
public:
   /// be schedule via the returned reference, i.e. only reads of the preamble
   /// are possible.
   std::shared_ptr<const ASTWorker> lock() { return Worker; }
 
 private:
   std::shared_ptr<ASTWorker> Worker;
 };
 
-ASTWorkerHandle ASTWorker::Create(llvm::StringRef File, AsyncTaskRunner *Tasks,
-                                  Semaphore &Barrier, CppFile AST,
-                                  steady_clock::duration UpdateDebounce) {
+ASTWorkerHandle ASTWorker::Create(PathRef FileName,
+                                  TUScheduler::ASTCache &IdleASTs,
+                                  AsyncTaskRunner *Tasks, Semaphore &Barrier,
+                                  steady_clock::duration UpdateDebounce,
+                                  std::shared_ptr<PCHContainerOperations> PCHs,
+                                  bool StorePreamblesInMemory,
+                                  PreambleParsedCallback PreambleCallback) {
   std::shared_ptr<ASTWorker> Worker(new ASTWorker(
-      File, Barrier, std::move(AST), /*RunSync=*/!Tasks, UpdateDebounce));
+      FileName, IdleASTs, Barrier, /*RunSync=*/!Tasks, UpdateDebounce,
+      std::move(PCHs), StorePreamblesInMemory, std::move(PreambleCallback)));
   if (Tasks)
-    Tasks->runAsync("worker:" + llvm::sys::path::filename(File),
+    Tasks->runAsync("worker:" + llvm::sys::path::filename(FileName),
                     [Worker]() { Worker->run(); });
 
   return ASTWorkerHandle(std::move(Worker));
 }
 
-ASTWorker::ASTWorker(llvm::StringRef File, Semaphore &Barrier, CppFile AST,
-                     bool RunSync, steady_clock::duration UpdateDebounce)
-    : File(File), RunSync(RunSync), UpdateDebounce(UpdateDebounce),
-      Barrier(Barrier), AST(std::move(AST)), Done(false) {
-  if (RunSync)
-    return;
-}
+ASTWorker::ASTWorker(PathRef FileName, TUScheduler::ASTCache &LRUCache,
+                     Semaphore &Barrier, bool RunSync,
+                     steady_clock::duration UpdateDebounce,
+                     std::shared_ptr<PCHContainerOperations> PCHs,
+                     bool StorePreamblesInMemory,
+                     PreambleParsedCallback PreambleCallback)
+    : IdleASTs(LRUCache), RunSync(RunSync), UpdateDebounce(UpdateDebounce),
+      FileName(FileName), StorePreambleInMemory(StorePreamblesInMemory),
+      PreambleCallback(std::move(PreambleCallback)), PCHs(std::move(PCHs)),
+      Barrier(Barrier), Done(false) {}
 
 ASTWorker::~ASTWorker() {
+  // Make sure we remove the cached AST, if any.
+  IdleASTs.take(this);
 #ifndef NDEBUG
   std::lock_guard<std::mutex> Lock(Mutex);
   assert(Done && "handle was not destroyed");
   assert(Requests.empty() && "unprocessed requests when destroying ASTWorker");
 #endif
 }
 
 void ASTWorker::update(ParseInputs Inputs, WantDiagnostics WantDiags,
                        UniqueFunction<void(std::vector<Diag>)> OnUpdated) {
   auto Task = [=](decltype(OnUpdated) OnUpdated) mutable {
+    tooling::CompileCommand OldCommand = std::move(FileInputs.CompileCommand);
     FileInputs = Inputs;
-    auto Diags = AST.rebuild(std::move(Inputs));
+    // Remove the old AST if it's still in cache.
+    IdleASTs.take(this);
+
+    log("Updating file " + FileName + " with command [" +
+        Inputs.CompileCommand.Directory + "] " +
+        llvm::join(Inputs.CompileCommand.CommandLine, " "));
+    // Rebuild the preamble and the AST.
+    Invocation = buildCompilerInvocation(Inputs);
+    if (!Invocation) {
+      log("Could not build CompilerInvocation for file " + FileName);
+      return;
+    }
 
+    std::shared_ptr<const PreambleData> NewPreamble = buildPreamble(
+        FileName, *Invocation, getPossiblyStalePreamble(), OldCommand, Inputs,
+        PCHs, StorePreambleInMemory, PreambleCallback);
     {
       std::lock_guard<std::mutex> Lock(Mutex);
-      if (AST.getPreamble())
-        LastBuiltPreamble = AST.getPreamble();
-      LastASTSize = AST.getUsedBytes();
+      if (NewPreamble)
+        LastBuiltPreamble = NewPreamble;
     }
+    // Build the AST for diagnostics.
+    llvm::Optional<ParsedAST> AST =
+        buildAST(FileName, llvm::make_unique<CompilerInvocation>(*Invocation),
+                 Inputs, NewPreamble, PCHs);
     // We want to report the diagnostics even if this update was cancelled.
     // It seems more useful than making the clients wait indefinitely if they
     // spam us with updates.
-    if (Diags && WantDiags != WantDiagnostics::No)
-      OnUpdated(std::move(*Diags));
+    if (WantDiags != WantDiagnostics::No && AST)
+      OnUpdated(AST->getDiagnostics());
+    // Stash the AST in the cache for further use.
+    IdleASTs.put(this,
+                 AST ? llvm::make_unique<ParsedAST>(std::move(*AST)) : nullptr);
   };
 
   startTask("Update", Bind(Task, std::move(OnUpdated)), WantDiags);
 }
 
 void ASTWorker::runWithAST(
     llvm::StringRef Name,
     UniqueFunction<void(llvm::Expected<InputsAndAST>)> Action) {
   auto Task = [=](decltype(Action) Action) {
-    ParsedAST *ActualAST = AST.getAST();
-    if (!ActualAST) {
-      Action(llvm::make_error<llvm::StringError>("invalid AST",
-                                                 llvm::errc::invalid_argument));
-      return;
-    }
-    Action(InputsAndAST{FileInputs, *ActualAST});
-
-    // Size of the AST might have changed after reads too, e.g. if some decls
-    // were deserialized from preamble.
-    std::lock_guard<std::mutex> Lock(Mutex);
-    LastASTSize = ActualAST->getUsedBytes();
+    llvm::Optional<std::unique_ptr<ParsedAST>> AST = IdleASTs.take(this);
+    if (!AST) {
+      // Try rebuilding the AST.
+      llvm::Optional<ParsedAST> NewAST =
+          Invocation
+              ? buildAST(FileName,
+                         llvm::make_unique<CompilerInvocation>(*Invocation),
+                         FileInputs, getPossiblyStalePreamble(), PCHs)
+              : llvm::None;
+      AST = NewAST ? llvm::make_unique<ParsedAST>(std::move(*NewAST)) : nullptr;
+    }
+    // Make sure we put the AST back into the LRU cache.
+    auto _ = llvm::make_scope_exit(
+        [&AST, this]() { IdleASTs.put(this, std::move(*AST)); });
+    // Run the user-provided action.
+    if (!*AST)
+      return Action(llvm::make_error<llvm::StringError>(
+          "invalid AST", llvm::errc::invalid_argument));
+    Action(InputsAndAST{FileInputs, **AST});
   };
-
   startTask(Name, Bind(Task, std::move(Action)),
             /*UpdateType=*/llvm::None);
 }
 
 std::shared_ptr<const PreambleData>
 ASTWorker::getPossiblyStalePreamble() const {
   std::lock_guard<std::mutex> Lock(Mutex);
   return LastBuiltPreamble;
 }
 
 std::size_t ASTWorker::getUsedBytes() const {
-  std::lock_guard<std::mutex> Lock(Mutex);
-  return LastASTSize;
+  // Note that we don't report the size of ASTs currently used for processing
+  // the in-flight requests. We used this information for debugging purposes
+  // only, so this should be fine.
+  std::size_t Result = IdleASTs.getUsedBytes(this);
+  if (auto Preamble = getPossiblyStalePreamble())
+    Result += Preamble->Preamble.getSize();
+  return Result;
 }
 
+bool ASTWorker::isASTCached() const { return IdleASTs.getUsedBytes(this) != 0; }
+
 void ASTWorker::stop() {
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     assert(!Done && "stop() called twice");
     Done = true;
   }
   RequestsCV.notify_all();
 }
 
 void ASTWorker::startTask(llvm::StringRef Name, UniqueFunction<void()> Task,
                           llvm::Optional<WantDiagnostics> UpdateType) {
   if (RunSync) {
     assert(!Done && "running a task after stop()");
-    trace::Span Tracer(Name + ":" + llvm::sys::path::filename(File));
+    trace::Span Tracer(Name + ":" + llvm::sys::path::filename(FileName));
     Task();
     return;
   }
 
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     assert(!Done && "running a task after stop()");
     Requests.push_back({std::move(Task), Name, steady_clock::now(),
struct TUScheduler::FileData {
   std::string Contents;
   tooling::CompileCommand Command;
   ASTWorkerHandle Worker;
 };
 
 TUScheduler::TUScheduler(unsigned AsyncThreadsCount,
                          bool StorePreamblesInMemory,
                          PreambleParsedCallback PreambleCallback,
-                         steady_clock::duration UpdateDebounce)
+                         std::chrono::steady_clock::duration UpdateDebounce,
+                         ASTRetentionPolicy RetentionPolicy)
     : StorePreamblesInMemory(StorePreamblesInMemory),
       PCHOps(std::make_shared<PCHContainerOperations>()),
       PreambleCallback(std::move(PreambleCallback)), Barrier(AsyncThreadsCount),
+      IdleASTs(llvm::make_unique<ASTCache>(RetentionPolicy.MaxRetainedASTs)),
       UpdateDebounce(UpdateDebounce) {
   if (0 < AsyncThreadsCount) {
     PreambleTasks.emplace();
     WorkerThreads.emplace();
   }
 }
 
 TUScheduler::~TUScheduler() {

 
 void TUScheduler::update(PathRef File, ParseInputs Inputs,
                          WantDiagnostics WantDiags,
                          UniqueFunction<void(std::vector<Diag>)> OnUpdated) {
   std::unique_ptr<FileData> &FD = Files[File];
   if (!FD) {
     // Create a new worker to process the AST-related tasks.
     ASTWorkerHandle Worker = ASTWorker::Create(
-        File, WorkerThreads ? WorkerThreads.getPointer() : nullptr, Barrier,
-        CppFile(File, StorePreamblesInMemory, PCHOps, PreambleCallback),
-        UpdateDebounce);
+        File, *IdleASTs, WorkerThreads ? WorkerThreads.getPointer() : nullptr,
+        Barrier, UpdateDebounce, PCHOps, StorePreamblesInMemory,
+        PreambleCallback);
     FD = std::unique_ptr<FileData>(new FileData{
         Inputs.Contents, Inputs.CompileCommand, std::move(Worker)});
   } else {
     FD->Contents = Inputs.Contents;
     FD->Command = Inputs.CompileCommand;
   }
   FD->Worker->update(std::move(Inputs), WantDiags, std::move(OnUpdated));
 }
TUScheduler::getUsedBytesPerFile() const {
   std::vector<std::pair<Path, std::size_t>> Result;
   Result.reserve(Files.size());
   for (auto &&PathAndFile : Files)
     Result.push_back(
         {PathAndFile.first(), PathAndFile.second->Worker->getUsedBytes()});
   return Result;
 }
 
+std::vector<Path> TUScheduler::getFilesWithCachedAST() const {
+  std::vector<Path> Result;
+  for (auto &&PathAndFile : Files) {
+    if (!PathAndFile.second->Worker->isASTCached())
+      continue;
+    Result.push_back(PathAndFile.first());
+  }
+  return Result;
+}
+
 } // namespace clangd
 } // namespace clang

clang-tools-extra/trunk/test/clangd/trace.test

 # RUN: env CLANGD_TRACE=%t clangd -lit-test < %s && FileCheck %s < %t
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
 ---
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{"uri":"test:///foo.c","languageId":"c","version":1,"text":"void main() {}"}}}
 # These assertions are a bit loose, to avoid brittleness.
 # CHECK: {"displayTimeUnit":"ns","traceEvents":[
 # CHECK: {
 # CHECK:   "args": {
 # CHECK:     "File": "{{.*(/|\\)}}foo.c"
 # CHECK:   },
-# CHECK:   "name": "Preamble",
+# CHECK:   "name": "BuildPreamble",
 # CHECK:   "ph": "X",
 # CHECK: }
 # CHECK: {
 # CHECK:   "args": {
 # CHECK:     "File": "{{.*(/|\\)}}foo.c"
 # CHECK:   },
-# CHECK:   "name": "Build",
+# CHECK:   "name": "BuildAST",
 # CHECK:   "ph": "X",
 # CHECK: }
 # CHECK: },
 ---
 {"jsonrpc":"2.0","id":5,"method":"shutdown"}

clang-tools-extra/trunk/unittests/clangd/FileIndexTests.cpp

 //===-- FileIndexTests.cpp  ---------------------------*- C++ -*-----------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "ClangdUnit.h"
 #include "TestFS.h"
 #include "TestTU.h"
 #include "index/FileIndex.h"
+#include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PCHContainerOperations.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 using testing::UnorderedElementsAre;
 
)cpp";
   });
   EXPECT_TRUE(SeenVector);
   EXPECT_TRUE(SeenMakeVector);
 }
 
 TEST(FileIndexTest, RebuildWithPreamble) {
   auto FooCpp = testPath("foo.cpp");
   auto FooH = testPath("foo.h");
-  FileIndex Index;
-  bool IndexUpdated = false;
-  CppFile File("foo.cpp", /*StorePreambleInMemory=*/true,
-               std::make_shared<PCHContainerOperations>(),
-               [&Index, &IndexUpdated](PathRef FilePath, ASTContext &Ctx,
-                                       std::shared_ptr<Preprocessor> PP) {
-                 EXPECT_FALSE(IndexUpdated)
-                     << "Expected only a single index update";
-                 IndexUpdated = true;
-                 Index.update(FilePath, &Ctx, std::move(PP));
-               });
-
   // Preparse ParseInputs.
   ParseInputs PI;
   PI.CompileCommand.Directory = testRoot();
   PI.CompileCommand.Filename = FooCpp;
   PI.CompileCommand.CommandLine = {"clang", "-xc++", FooCpp};
 
   llvm::StringMap<std::string> Files;
   Files[FooCpp] = "";
   Files[FooH] = R"cpp(
     namespace ns_in_header {
       int func_in_header();
     }
   )cpp";
   PI.FS = buildTestFS(std::move(Files));
 
   PI.Contents = R"cpp(
     #include "foo.h"
     namespace ns_in_source {
       int func_in_source();
     }
   )cpp";
 
   // Rebuild the file.
-  File.rebuild(std::move(PI));
+  auto CI = buildCompilerInvocation(PI);
+
+  FileIndex Index;
+  bool IndexUpdated = false;
+  buildPreamble(
+      FooCpp, *CI, /*OldPreamble=*/nullptr, tooling::CompileCommand(), PI,
+      std::make_shared<PCHContainerOperations>(), /*StoreInMemory=*/true,
+      [&Index, &IndexUpdated](PathRef FilePath, ASTContext &Ctx,
+                              std::shared_ptr<Preprocessor> PP) {
+        EXPECT_FALSE(IndexUpdated) << "Expected only a single index update";
+        IndexUpdated = true;
+        Index.update(FilePath, &Ctx, std::move(PP));
+      });
   ASSERT_TRUE(IndexUpdated);
 
   // Check the index contains symbols from the preamble, but not from the main
   // file.
   FuzzyFindRequest Req;
   Req.Query = "";
   Req.Scopes = {"", "ns_in_header::"};
 
   EXPECT_THAT(
       match(Index, Req),
       UnorderedElementsAre("ns_in_header", "ns_in_header::func_in_header"));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

clang-tools-extra/trunk/unittests/clangd/TUSchedulerTests.cpp

 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <algorithm>
 #include <utility>
 
 namespace clang {
 namespace clangd {
 
+using ::testing::_;
+using ::testing::AnyOf;
 using ::testing::Pair;
 using ::testing::Pointee;
+using ::testing::UnorderedElementsAre;
 
 void ignoreUpdate(llvm::Optional<std::vector<Diag>>) {}
 void ignoreError(llvm::Error Err) {
   handleAllErrors(std::move(Err), [](const llvm::ErrorInfoBase &) {});
 }
 
 class TUSchedulerTests : public ::testing::Test {
 protected:
   ParseInputs getInputs(PathRef File, std::string Contents) {
     return ParseInputs{*CDB.getCompileCommand(File), buildTestFS(Files),
                        std::move(Contents)};
   }
 
   llvm::StringMap<std::string> Files;
 
 private:
   MockCompilationDatabase CDB;
 };
 
 TEST_F(TUSchedulerTests, MissingFiles) {
   TUScheduler S(getDefaultAsyncThreadsCount(),
                 /*StorePreamblesInMemory=*/true,
                 /*PreambleParsedCallback=*/nullptr,
-                /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero());
+                /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero(),
+                ASTRetentionPolicy());
 
   auto Added = testPath("added.cpp");
   Files[Added] = "";
 
   auto Missing = testPath("missing.cpp");
   Files[Missing] = "";
 
   S.update(Added, getInputs(Added, ""), WantDiagnostics::No, ignoreUpdate);
TEST_F(TUSchedulerTests, WantDiagnostics) {
   {
     // To avoid a racy test, don't allow tasks to actualy run on the worker
     // thread until we've scheduled them all.
     Notification Ready;
     TUScheduler S(
         getDefaultAsyncThreadsCount(),
         /*StorePreamblesInMemory=*/true,
         /*PreambleParsedCallback=*/nullptr,
-        /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero());
+        /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero(),
+        ASTRetentionPolicy());
     auto Path = testPath("foo.cpp");
     S.update(Path, getInputs(Path, ""), WantDiagnostics::Yes,
              [&](std::vector<Diag>) { Ready.wait(); });
 
     S.update(Path, getInputs(Path, "request diags"), WantDiagnostics::Yes,
              [&](std::vector<Diag> Diags) { ++CallbackCount; });
     S.update(Path, getInputs(Path, "auto (clobbered)"), WantDiagnostics::Auto,
              [&](std::vector<Diag> Diags) {

 }
 
 TEST_F(TUSchedulerTests, Debounce) {
   std::atomic<int> CallbackCount(0);
   {
     TUScheduler S(getDefaultAsyncThreadsCount(),
                   /*StorePreamblesInMemory=*/true,
                   /*PreambleParsedCallback=*/nullptr,
-                  /*UpdateDebounce=*/std::chrono::seconds(1));
+                  /*UpdateDebounce=*/std::chrono::seconds(1),
+                  ASTRetentionPolicy());
     // FIXME: we could probably use timeouts lower than 1 second here.
     auto Path = testPath("foo.cpp");
     S.update(Path, getInputs(Path, "auto (debounced)"), WantDiagnostics::Auto,
              [&](std::vector<Diag> Diags) {
                ADD_FAILURE() << "auto should have been debounced and canceled";
              });
     std::this_thread::sleep_for(std::chrono::milliseconds(200));
     S.update(Path, getInputs(Path, "auto (timed out)"), WantDiagnostics::Auto,
TEST_F(TUSchedulerTests, ManyUpdates) {
   int TotalPreambleReads = 0;
   int TotalUpdates = 0;
 
   // Run TUScheduler and collect some stats.
   {
     TUScheduler S(getDefaultAsyncThreadsCount(),
                   /*StorePreamblesInMemory=*/true,
                   /*PreambleParsedCallback=*/nullptr,
-                  /*UpdateDebounce=*/std::chrono::milliseconds(50));
+                  /*UpdateDebounce=*/std::chrono::milliseconds(50),
+                  ASTRetentionPolicy());
 
     std::vector<std::string> Files;
     for (int I = 0; I < FilesCount; ++I) {
       std::string Name = "foo" + std::to_string(I) + ".cpp";
       Files.push_back(testPath(Name));
       this->Files[Files.back()] = "";
     }
 
for (int FileI = 0; FileI < FilesCount; ++FileI) {
 
                          std::lock_guard<std::mutex> Lock(Mut);
                          ++TotalASTReads;
                        });
         }
 
         {
           WithContextValue WithNonce(NonceKey, ++Nonce);
-          S.runWithPreamble(
-              "CheckPreamble", File,
+          S.runWithPreamble("CheckPreamble", File,
                             [Inputs, Nonce, &Mut, &TotalPreambleReads](
                                 llvm::Expected<InputsAndPreamble> Preamble) {
-                EXPECT_THAT(Context::current().get(NonceKey), Pointee(Nonce));
+                              EXPECT_THAT(Context::current().get(NonceKey),
+                                          Pointee(Nonce));
 
                               ASSERT_TRUE((bool)Preamble);
                               EXPECT_EQ(Preamble->Contents, Inputs.Contents);
 
                               std::lock_guard<std::mutex> Lock(Mut);
                               ++TotalPreambleReads;
                             });
         }
       }
     }
   } // TUScheduler destructor waits for all operations to finish.
 
   std::lock_guard<std::mutex> Lock(Mut);
   EXPECT_EQ(TotalUpdates, FilesCount * UpdatesPerFile);
   EXPECT_EQ(TotalASTReads, FilesCount * UpdatesPerFile);
   EXPECT_EQ(TotalPreambleReads, FilesCount * UpdatesPerFile);
 }
 
+TEST_F(TUSchedulerTests, EvictedAST) {
+  ASTRetentionPolicy Policy;
+  Policy.MaxRetainedASTs = 2;
+  TUScheduler S(
+      /*AsyncThreadsCount=*/1, /*StorePreambleInMemory=*/true,
+      PreambleParsedCallback(),
+      /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero(), Policy);
+
+  llvm::StringLiteral SourceContents = R"cpp(
+    int* a;
+    double* b = a;
+  )cpp";
+
+  auto Foo = testPath("foo.cpp");
+  auto Bar = testPath("bar.cpp");
+  auto Baz = testPath("baz.cpp");
+
+  std::atomic<int> BuiltASTCounter;
+  BuiltASTCounter = false;
+  // Build one file in advance. We will not access it later, so it will be the
+  // one that the cache will evict.
+  S.update(Foo, getInputs(Foo, SourceContents), WantDiagnostics::Yes,
+           [&BuiltASTCounter](std::vector<Diag> Diags) { ++BuiltASTCounter; });
+  ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(1)));
+  ASSERT_EQ(BuiltASTCounter.load(), 1);
+
+  // Build two more files. Since we can retain only 2 ASTs, these should be the
+  // ones we see in the cache later.
+  S.update(Bar, getInputs(Bar, SourceContents), WantDiagnostics::Yes,
+           [&BuiltASTCounter](std::vector<Diag> Diags) { ++BuiltASTCounter; });
+  S.update(Baz, getInputs(Baz, SourceContents), WantDiagnostics::Yes,
+           [&BuiltASTCounter](std::vector<Diag> Diags) { ++BuiltASTCounter; });
+  ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(1)));
+  ASSERT_EQ(BuiltASTCounter.load(), 3);
+
+  // Check only the last two ASTs are retained.
+  ASSERT_THAT(S.getFilesWithCachedAST(), UnorderedElementsAre(Bar, Baz));
+
+  // Access the old file again.
+  S.update(Foo, getInputs(Foo, SourceContents), WantDiagnostics::Yes,
+           [&BuiltASTCounter](std::vector<Diag> Diags) { ++BuiltASTCounter; });
+  ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(1)));
+  ASSERT_EQ(BuiltASTCounter.load(), 4);
+
+  // Check the AST for foo.cpp is retained now and one of the others got
+  // evicted.
+  EXPECT_THAT(S.getFilesWithCachedAST(),
+              UnorderedElementsAre(Foo, AnyOf(Bar, Baz)));
+}
+
 } // namespace clangd
 } // namespace clang