[clangd] Support external throttler for preamble builds

Building preambles is the most resource-intensive thing clangd does, driving
peak RAM and sustained CPU usage.

In a hosted environment where multiple clangd instances are packed into the same
container, it's useful to be able to limit the *aggregate* resource peaks.

Differential Revision: https://reviews.llvm.org/D129100

clang-tools-extra/clangd/ClangdServer.cpp

@@ -166,6 +166,7 @@ ClangdServer::Options::operator TUScheduler::Options() const {
   Opts.StorePreamblesInMemory = StorePreamblesInMemory;
   Opts.UpdateDebounce = UpdateDebounce;
   Opts.ContextProvider = ContextProvider;
+  Opts.PreambleThrottler = PreambleThrottler;
   return Opts;
 }
 

clang-tools-extra/clangd/ClangdServer.h

@@ -104,6 +104,9 @@ class ClangdServer {
     /// Cached preambles are potentially large. If false, store them on disk.
     bool StorePreamblesInMemory = true;
 
+    /// This throttler controls which preambles may be built at a given time.
+    clangd::PreambleThrottler *PreambleThrottler = nullptr;
+
     /// 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;

clang-tools-extra/clangd/TUScheduler.cpp

@@ -381,6 +381,41 @@ class SynchronizedTUStatus {
   ParsingCallbacks &Callbacks;
 };
 
+// An attempt to acquire resources for a task using PreambleThrottler.
+// Initially it is unsatisfied, it (hopefully) becomes satisfied later but may
+// be destroyed before then. Destruction releases all resources.
+class PreambleThrottlerRequest {
+public:
+  // The condition variable is signalled when the request is satisfied.
+  PreambleThrottlerRequest(llvm::StringRef Filename,
+                           PreambleThrottler *Throttler,
+                           std::condition_variable &CV)
+      : Throttler(Throttler), Satisfied(Throttler == nullptr) {
+    // If there is no throttler, this dummy request is always satisfied.
+    if (!Throttler)
+      return;
+    ID = Throttler->acquire(Filename, [&] {
+      Satisfied.store(true, std::memory_order_release);
+      CV.notify_all();
+    });
+  }
+
+  bool satisfied() const { return Satisfied.load(std::memory_order_acquire); }
+
+  // When the request is destroyed:
+  //  - if resources are not yet obtained, stop trying to get them.
+  //  - if resources were obtained, release them.
+  ~PreambleThrottlerRequest() {
+    if (Throttler)
+      Throttler->release(ID);
+  }
+
+private:
+  PreambleThrottler::RequestID ID;
+  PreambleThrottler *Throttler;
+  std::atomic<bool> Satisfied = {false};
+};
+
 /// Responsible for building preambles. Whenever the thread is idle and the
 /// preamble is outdated, it starts to build a fresh preamble from the latest
 /// inputs. If RunSync is true, preambles are built synchronously in update()
@@ -389,12 +424,13 @@ class PreambleThread {
 public:
   PreambleThread(llvm::StringRef FileName, ParsingCallbacks &Callbacks,
                  bool StorePreambleInMemory, bool RunSync,
-                 SynchronizedTUStatus &Status,
+                 PreambleThrottler *Throttler, SynchronizedTUStatus &Status,
                  TUScheduler::HeaderIncluderCache &HeaderIncluders,
                  ASTWorker &AW)
       : FileName(FileName), Callbacks(Callbacks),
-        StoreInMemory(StorePreambleInMemory), RunSync(RunSync), Status(Status),
-        ASTPeer(AW), HeaderIncluders(HeaderIncluders) {}
+        StoreInMemory(StorePreambleInMemory), RunSync(RunSync),
+        Throttler(Throttler), Status(Status), ASTPeer(AW),
+        HeaderIncluders(HeaderIncluders) {}
 
   /// It isn't guaranteed that each requested version will be built. If there
   /// are multiple update requests while building a preamble, only the last one
@@ -428,13 +464,27 @@ class PreambleThread {
 
   void run() {
     while (true) {
+      llvm::Optional<PreambleThrottlerRequest> Throttle;
       {
         std::unique_lock<std::mutex> Lock(Mutex);
         assert(!CurrentReq && "Already processing a request?");
         // Wait until stop is called or there is a request.
-        ReqCV.wait(Lock, [this] { return NextReq || Done; });
+        ReqCV.wait(Lock, [&] { return NextReq || Done; });
+        if (Done)
+          break;
+
+        Throttle.emplace(FileName, Throttler, ReqCV);
+        // If acquire succeeded synchronously, avoid status jitter.
+        if (!Throttle->satisfied())
+          Status.update([&](TUStatus &Status) {
+            Status.PreambleActivity = PreambleAction::Queued;
+          });
+        ReqCV.wait(Lock, [&] { return Throttle->satisfied() || Done; });
         if (Done)
           break;
+        // While waiting for the throttler, the request may have been updated!
+        // That's fine though, there's still guaranteed to be some request.
+
         CurrentReq = std::move(*NextReq);
         NextReq.reset();
       }
@@ -518,6 +568,7 @@ class PreambleThread {
   ParsingCallbacks &Callbacks;
   const bool StoreInMemory;
   const bool RunSync;
+  PreambleThrottler *Throttler;
 
   SynchronizedTUStatus &Status;
   ASTWorker &ASTPeer;
@@ -778,7 +829,7 @@ ASTWorker::ASTWorker(PathRef FileName, const GlobalCompilationDatabase &CDB,
       ContextProvider(Opts.ContextProvider), CDB(CDB), Callbacks(Callbacks),
       Barrier(Barrier), Done(false), Status(FileName, Callbacks),
       PreamblePeer(FileName, Callbacks, Opts.StorePreamblesInMemory, RunSync,
-                   Status, HeaderIncluders, *this) {
+                   Opts.PreambleThrottler, Status, HeaderIncluders, *this) {
   // Set a fallback command because compile command can be accessed before
   // `Inputs` is initialized. Other fields are only used after initialization
   // from client inputs.
@@ -1499,6 +1550,9 @@ std::string renderTUAction(const PreambleAction PA, const ASTAction &AA) {
   case PreambleAction::Building:
     Result.push_back("parsing includes");
     break;
+  case PreambleAction::Queued:
+    Result.push_back("includes are queued");
+    break;
   case PreambleAction::Idle:
     // We handle idle specially below.
     break;

clang-tools-extra/clangd/TUScheduler.h

@@ -87,9 +87,38 @@ struct DebouncePolicy {
   static DebouncePolicy fixed(clock::duration);
 };
 
+/// PreambleThrottler controls which preambles can build at any given time.
+/// This can be used to limit overall concurrency, and to prioritize some
+/// preambles over others.
+/// In a distributed environment, a throttler may be able to coordinate resource
+/// use across several clangd instances.
+///
+/// This class is threadsafe.
+class PreambleThrottler {
+public:
+  virtual ~PreambleThrottler() = default;
+
+  using RequestID = unsigned;
+  using Callback = llvm::unique_function<void()>;
+  /// Attempt to acquire resources to build a file's preamble.
+  ///
+  /// Does not block, may eventually invoke the callback to satisfy the request.
+  /// If the callback is invoked, release() must be called afterwards.
+  virtual RequestID acquire(llvm::StringRef Filename, Callback);
+  /// Abandons the request/releases any resources that have been acquired.
+  ///
+  /// Must be called exactly once after acquire().
+  /// acquire()'s callback will not be invoked after release() returns.
+  virtual void release(RequestID) = 0;
+
+  // FIXME: we may want to be able attach signals to filenames.
+  //        this would allow the throttler to make better scheduling decisions.
+};
+
 enum class PreambleAction {
-  Idle,
+  Queued,
   Building,
+  Idle,
 };
 
 struct ASTAction {
@@ -200,6 +229,9 @@ class TUScheduler {
     /// Determines when to keep idle ASTs in memory for future use.
     ASTRetentionPolicy RetentionPolicy;
 
+    /// This throttler controls which preambles may be built at a given time.
+    clangd::PreambleThrottler *PreambleThrottler = nullptr;
+
     /// Used to create a context that wraps each single operation.
     /// Typically to inject per-file configuration.
     /// If the path is empty, context sholud be "generic".

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

@@ -1372,6 +1372,150 @@ TEST_F(TUSchedulerTests, PreservesLastActiveFile) {
     CheckNoFileActionsSeesLastActiveFile(LastActive);
   }
 }
+
+TEST_F(TUSchedulerTests, PreambleThrottle) {
+  const int NumRequests = 4;
+  // Silly throttler that waits for 4 requests, and services them in reverse.
+  // Doesn't honor cancellation but records it.
+  struct : public PreambleThrottler {
+    std::mutex Mu;
+    std::vector<std::string> Acquires;
+    std::vector<RequestID> Releases;
+    llvm::DenseMap<RequestID, Callback> Callbacks;
+    // If set, the notification is signalled after acquiring the specified ID.
+    llvm::Optional<std::pair<RequestID, Notification *>> Notify;
+
+    RequestID acquire(llvm::StringRef Filename, Callback CB) override {
+      RequestID ID;
+      Callback Invoke;
+      {
+        std::lock_guard<std::mutex> Lock(Mu);
+        ID = Acquires.size();
+        Acquires.emplace_back(Filename);
+        // If we're full, satisfy this request immediately.
+        if (Acquires.size() == NumRequests) {
+          Invoke = std::move(CB);
+        } else {
+          Callbacks.try_emplace(ID, std::move(CB));
+        }
+      }
+      if (Invoke)
+        Invoke();
+      if (Notify && ID == Notify->first) {
+        Notify->second->notify();
+        Notify.reset();
+      }
+      return ID;
+    }
+
+    void release(RequestID ID) override {
+      Releases.push_back(ID);
+      Callback SatisfyNext;
+      {
+        std::lock_guard<std::mutex> Lock(Mu);
+        if (ID > 0)
+          SatisfyNext = std::move(Callbacks[ID - 1]);
+      }
+      if (SatisfyNext)
+        SatisfyNext();
+    }
+
+    void reset() {
+      Acquires.clear();
+      Releases.clear();
+      Callbacks.clear();
+    }
+  } Throttler;
+
+  struct CaptureBuiltFilenames : public ParsingCallbacks {
+    std::vector<std::string> &Filenames;
+    CaptureBuiltFilenames(std::vector<std::string> &Filenames)
+        : Filenames(Filenames) {}
+    void onPreambleAST(PathRef Path, llvm::StringRef Version,
+                       const CompilerInvocation &CI, ASTContext &Ctx,
+                       Preprocessor &PP, const CanonicalIncludes &) override {
+      // Deliberately no synchronization.
+      // The PreambleThrottler should serialize these calls, if not then tsan
+      // will find a bug here.
+      Filenames.emplace_back(Path);
+    }
+  };
+
+  auto Opts = optsForTest();
+  Opts.AsyncThreadsCount = 2 * NumRequests; // throttler is the bottleneck
+  Opts.PreambleThrottler = &Throttler;
+
+  std::vector<std::string> Filenames;
+
+  {
+    std::vector<std::string> BuiltFilenames;
+    TUScheduler S(CDB, Opts,
+                  std::make_unique<CaptureBuiltFilenames>(BuiltFilenames));
+    for (unsigned I = 0; I < NumRequests; ++I) {
+      auto Path = testPath(std::to_string(I) + ".cc");
+      Filenames.push_back(Path);
+      S.update(Path, getInputs(Path, ""), WantDiagnostics::Yes);
+    }
+    ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
+
+    // The throttler saw all files, and we built them.
+    EXPECT_THAT(Throttler.Acquires,
+                testing::UnorderedElementsAreArray(Filenames));
+    EXPECT_THAT(BuiltFilenames,
+                testing::UnorderedElementsAreArray(Filenames));
+    // We built the files in reverse order that the throttler saw them.
+    EXPECT_THAT(BuiltFilenames,
+                testing::ElementsAreArray(Throttler.Acquires.rbegin(),
+                                          Throttler.Acquires.rend()));
+    // Resources for each file were correctly released.
+    EXPECT_THAT(Throttler.Releases, ElementsAre(3, 2, 1, 0));
+  }
+
+  Throttler.reset();
+
+  // This time, enqueue 2 files, then cancel one of them while still waiting.
+  // Finally shut down the server. Observe that everything gets cleaned up.
+  Notification AfterAcquire2;
+  Notification AfterFinishA;
+  Throttler.Notify = {1, &AfterAcquire2};
+  std::vector<std::string> BuiltFilenames;
+  auto A = testPath("a.cc");
+  auto B = testPath("b.cc");
+  Filenames = {A, B};
+  {
+    TUScheduler S(CDB, Opts,
+                  std::make_unique<CaptureBuiltFilenames>(BuiltFilenames));
+    updateWithCallback(S, A, getInputs(A, ""), WantDiagnostics::Yes,
+                       [&] { AfterFinishA.notify(); });
+    S.update(B, getInputs(B, ""), WantDiagnostics::Yes);
+    AfterAcquire2.wait();
+
+    // The throttler saw all files, but we built none.
+    EXPECT_THAT(Throttler.Acquires,
+                testing::UnorderedElementsAreArray(Filenames));
+    EXPECT_THAT(BuiltFilenames, testing::IsEmpty());
+    // We haven't released anything yet, we're still waiting.
+    EXPECT_THAT(Throttler.Releases, testing::IsEmpty());
+
+    // Now close file A, which will shut down its AST worker.
+    S.remove(A);
+    // Request is destroyed after the queue shutdown, so release() has happened.
+    AfterFinishA.wait();
+    // We still didn't build anything.
+    EXPECT_THAT(BuiltFilenames, testing::IsEmpty());
+    // But we've cancelled the request to build A (not sure which its ID is).
+    EXPECT_THAT(Throttler.Releases, ElementsAre(AnyOf(1, 0)));
+
+    // Now shut down the TU Scheduler.
+  }
+  // The throttler saw all files, but we built none.
+  EXPECT_THAT(Throttler.Acquires,
+              testing::UnorderedElementsAreArray(Filenames));
+  EXPECT_THAT(BuiltFilenames, testing::IsEmpty());
+  // We gave up waiting and everything got released (in some order).
+  EXPECT_THAT(Throttler.Releases, UnorderedElementsAre(1, 0));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang