[clangd] Use atomics instead of locks to track periodic memory trimming

[sammccall]

Summary

Instead of always locking/unlocking a contended mutex, we now do one atomic read
in the common case, and one read + one exchange if the timer has expried.

Also use this for memory profiling which has similar/compatible requirements.

[qchateau]

LGTM

Small nits:

• operator() is not const but Next is mutable, seems to me you intended to have operator() as const
• memory orders for the atomic operations can probably be downgraded to std::memory_order_acquire/std::memory_order_acq_rel. I think the first load can even be relaxed but that I'm always careful with these

[sammccall]

In D93726#2468855, @qchateau wrote:

LGTM

Small nits:

• operator() is not const but Next is mutable, seems to me you intended to have operator() as const

Oops, I originally did plan to make operator() const, but then decided it'd be clearer for the caller to make PeriodicThrottler mutable instead. (Because pretending it doesn't have state is confusing).
Removed mutable from Next.

• memory orders for the atomic operations can probably be downgraded to std::memory_order_acquire/std::memory_order_acq_rel. I think the first load can even be relaxed but that I'm always careful with these

Done. I have to confess I often pretend memory orders other than seq_cst don't exist just to keep the cognitive load down, but this case seems clear/isolated enough.

clang-tools-extra/clangd/ClangdLSPServer.h

 #include "Features.inc"
 #include "FindSymbols.h"
 #include "GlobalCompilationDatabase.h"
 #include "Protocol.h"
 #include "Transport.h"
 #include "support/Context.h"
 #include "support/MemoryTree.h"
 #include "support/Path.h"
+#include "support/Threading.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/JSON.h"
 #include <chrono>
 #include <cstddef>
 #include <memory>
 
private:
   publishTheiaSemanticHighlighting(const TheiaSemanticHighlightingParams &);
 
   /// Sends a "publishDiagnostics" notification to the LSP client.
   void publishDiagnostics(const PublishDiagnosticsParams &);
 
   /// Runs profiling and exports memory usage metrics if tracing is enabled and
   /// profiling hasn't happened recently.
   void maybeExportMemoryProfile();
+  PeriodicThrottler ShouldProfile;
 
   /// Run the MemoryCleanup callback if it's time.
   /// This method is thread safe.
   void maybeCleanupMemory();
-
-  /// Timepoint until which profiling is off. It is used to throttle profiling
-  /// requests.
-  std::chrono::steady_clock::time_point NextProfileTime;
-
-  /// Next time we want to call the MemoryCleanup callback.
-  std::mutex NextMemoryCleanupTimeMutex;
-  std::chrono::steady_clock::time_point NextMemoryCleanupTime;
+  PeriodicThrottler ShouldCleanupMemory;
 
   /// Since initialization of CDBs and ClangdServer is done lazily, the
   /// following context captures the one used while creating ClangdLSPServer and
   /// passes it to above mentioned object instances to make sure they share the
   /// same state.
   Context BackgroundContext;
 
   /// Used to indicate that the 'shutdown' request was received from the

clang-tools-extra/clangd/ClangdLSPServer.cpp

 }
 
 void ClangdLSPServer::publishDiagnostics(
     const PublishDiagnosticsParams &Params) {
   notify("textDocument/publishDiagnostics", Params);
 }
 
 void ClangdLSPServer::maybeExportMemoryProfile() {
-  if (!trace::enabled())
-    return;
-  // Profiling might be expensive, so we throttle it to happen once every 5
-  // minutes.
-  static constexpr auto ProfileInterval = std::chrono::minutes(5);
-  auto Now = std::chrono::steady_clock::now();
-  if (Now < NextProfileTime)
+  if (!trace::enabled() || !ShouldProfile())
     return;
 
   static constexpr trace::Metric MemoryUsage(
       "memory_usage", trace::Metric::Value, "component_name");
   trace::Span Tracer("ProfileBrief");
   MemoryTree MT;
   profile(MT);
   record(MT, "clangd_lsp_server", MemoryUsage);
-  NextProfileTime = Now + ProfileInterval;
 }
 
 void ClangdLSPServer::maybeCleanupMemory() {
-  // Memory cleanup is probably expensive, throttle it
-  static constexpr auto MemoryCleanupInterval = std::chrono::minutes(1);
-
-  if (!Opts.MemoryCleanup)
-    return;
-
-  // FIXME: this can probably be done without a mutex
-  // and the logic could be shared with maybeExportMemoryProfile
-  {
-    auto Now = std::chrono::steady_clock::now();
-    std::lock_guard<std::mutex> Lock(NextMemoryCleanupTimeMutex);
-    if (Now < NextMemoryCleanupTime)
-      return;
-    NextMemoryCleanupTime = Now + MemoryCleanupInterval;
-  }
-
-  vlog("Calling memory cleanup callback");
+  if (!Opts.MemoryCleanup || !ShouldCleanupMemory())
   Opts.MemoryCleanup();
 }
 
 // FIXME: This function needs to be properly tested.
 void ClangdLSPServer::onChangeConfiguration(
     const DidChangeConfigurationParams &Params) {
   applyConfiguration(Params.settings);
 }

 void ClangdLSPServer::onAST(const ASTParams &Params,
                             Callback<llvm::Optional<ASTNode>> CB) {
   Server->getAST(Params.textDocument.uri.file(), Params.range, std::move(CB));
 }
 
 ClangdLSPServer::ClangdLSPServer(class Transport &Transp,
                                  const ThreadsafeFS &TFS,
                                  const ClangdLSPServer::Options &Opts)
-    : BackgroundContext(Context::current().clone()), Transp(Transp),
+    : ShouldProfile(/*Period=*/std::chrono::minutes(5),
+                    /*Delay=*/std::chrono::minutes(1)),
+      ShouldCleanupMemory(/*Period=*/std::chrono::minutes(1),
+                          /*Delay=*/std::chrono::minutes(1)),
+      BackgroundContext(Context::current().clone()), Transp(Transp),
       MsgHandler(new MessageHandler(*this)), TFS(TFS),
       SupportedSymbolKinds(defaultSymbolKinds()),
       SupportedCompletionItemKinds(defaultCompletionItemKinds()), Opts(Opts) {
+
   // clang-format off
   MsgHandler->bind("initialize", &ClangdLSPServer::onInitialize);
   MsgHandler->bind("initialized", &ClangdLSPServer::onInitialized);
   MsgHandler->bind("shutdown", &ClangdLSPServer::onShutdown);
   MsgHandler->bind("sync", &ClangdLSPServer::onSync);
   MsgHandler->bind("textDocument/rangeFormatting", &ClangdLSPServer::onDocumentRangeFormatting);
   MsgHandler->bind("textDocument/onTypeFormatting", &ClangdLSPServer::onDocumentOnTypeFormatting);
   MsgHandler->bind("textDocument/formatting", &ClangdLSPServer::onDocumentFormatting);
ClangdLSPServer::ClangdLSPServer(class Transport &Transp,
   MsgHandler->bind("textDocument/selectionRange", &ClangdLSPServer::onSelectionRange);
   MsgHandler->bind("textDocument/documentLink", &ClangdLSPServer::onDocumentLink);
   MsgHandler->bind("textDocument/semanticTokens/full", &ClangdLSPServer::onSemanticTokens);
   MsgHandler->bind("textDocument/semanticTokens/full/delta", &ClangdLSPServer::onSemanticTokensDelta);
   MsgHandler->bind("$/memoryUsage", &ClangdLSPServer::onMemoryUsage);
   if (Opts.FoldingRanges)
     MsgHandler->bind("textDocument/foldingRange", &ClangdLSPServer::onFoldingRange);
   // clang-format on
-
-  // Delay first profile and memory cleanup until we've finished warming up.
-  NextMemoryCleanupTime = NextProfileTime =
-      std::chrono::steady_clock::now() + std::chrono::minutes(1);
 }
 
 ClangdLSPServer::~ClangdLSPServer() {
   IsBeingDestroyed = true;
   // Explicitly destroy ClangdServer first, blocking on threads it owns.
   // This ensures they don't access any other members.
   Server.reset();
 }

clang-tools-extra/clangd/support/Threading.h

auto V = Compute();
       // Insert into cache may fail if we raced with another thread.
       if (!R.second)
         return R.first->second; // Canonical value, from other thread.
     }
     return V;
   }
 };
 
+/// Used to guard an operation that should run at most every N seconds.
+///
+/// Usage:
+///   mutable PeriodicThrottler ShouldLog(std::chrono::seconds(1));
+///   void calledFrequently() {
+///     if (ShouldLog())
+///       log("this is not spammy");
+///   }
+///
+/// This class is threadsafe. If multiple threads are involved, then the guarded
+/// operation still needs to be threadsafe!
+class PeriodicThrottler {
+  using Stopwatch = std::chrono::steady_clock;
+  using Rep = Stopwatch::duration::rep;
+
+  Rep Period;
+  mutable std::atomic<Rep> Next;
+
+public:
+  /// If Period is zero, the throttler will return true every time.
+  PeriodicThrottler(Stopwatch::duration Period, Stopwatch::duration Delay = {})
+      : Period(Period.count()),
+        Next((Stopwatch::now() + Delay).time_since_epoch().count()) {}
+
+  /// Returns whether the operation should run at this time.
+  /// operator() is safe to call concurrently.
+  bool operator()();
+};
+
 } // namespace clangd
 } // namespace clang
 #endif

clang-tools-extra/clangd/support/Threading.cpp

void wait(std::unique_lock<std::mutex> &Lock, std::condition_variable &CV,
           Deadline D) {
   if (D == Deadline::zero())
     return;
   if (D == Deadline::infinity())
     return CV.wait(Lock);
   CV.wait_until(Lock, D.time());
 }
 
+bool PeriodicThrottler::operator()() {
+  Rep Now = Stopwatch::now().time_since_epoch().count();
+  Rep OldNext = Next.load();
+  if (Now < OldNext)
+    return false;
+  // We're ready to run (but may be racing other threads).
+  // Work out the updated target time, and run if we successfully bump it.
+  Rep NewNext = Now + Period;
+  return Next.compare_exchange_strong(OldNext, NewNext);
+}
+
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/unittests/support/ThreadingTests.cpp

 //===-- ThreadingTests.cpp --------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "support/Threading.h"
 #include "llvm/ADT/DenseMap.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include <chrono>
 #include <mutex>
 
 namespace clang {
 namespace clangd {
 class ThreadingTest : public ::testing::Test {};
 
 TEST_F(ThreadingTest, TaskRunner) {
   const int TasksCnt = 100;
Tasks.runAsync("A", [&] {
     });
   });
   Tasks.wait();
 
   ASSERT_EQ(ValueA, ValueB);
   ASSERT_THAT(ValueA.load(), testing::AnyOf('A', 'B'));
 }
 
+// It's hard to write a real test of this class, std::chrono is awkward to mock.
+// But test some degenerate cases at least.
+TEST(PeriodicThrottlerTest, Minimal) {
+  PeriodicThrottler Once(std::chrono::hours(24));
+  EXPECT_TRUE(Once());
+  EXPECT_FALSE(Once());
+  EXPECT_FALSE(Once());
+
+  PeriodicThrottler Later(std::chrono::hours(24),
+                          /*Delay=*/std::chrono::hours(24));
+  EXPECT_FALSE(Later());
+  EXPECT_FALSE(Later());
+  EXPECT_FALSE(Later());
+
+  PeriodicThrottler Always(std::chrono::seconds(0));
+  EXPECT_TRUE(Always());
+  EXPECT_TRUE(Always());
+  EXPECT_TRUE(Always());
+}
+
 } // namespace clangd
 } // namespace clang