[clangd] Propogate context in LSPServer tests

[kadircet]

Summary

This enables users to wait until async request completes. Taking a
response isn't enough, as receiving a response doesn't imply destruction of the
reqeust context.

This also fixes the raciness in lspserver latency metric test. As it needs to
wait for context destruction, rather than receiving a reply. Happy to perform
that fix in CallResult instead, i.e change LSPClient::call to perform the
insertion of Notification into context and change CallResult::take to wait for
destruction of context rather than receiving a reply(calling set).

[sammccall]

Hmm, I wonder if we should have a Context.bind(function) -> function.
I guess it runs into the usual thing of not being able to deduce a functor's signature, so you get an ugly templated return type..

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

120

which waiters? isn't it just this thread?

122–125

nit: scope WithContext to exclude taking the lock?

[kadircet]

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

120

i was envisioning the future(thought this already had a blockUntilIdle :D)

[kadircet]

• Address comments

[sammccall]

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

164

Sorry, I didn't really put all the pieces together in my head the first time around.
The context propagation seems OK, but it's too fiddly as a way to control sequencing - can't you just call Client.sync() and assert after that?

[kadircet]

• Call sync instead of waiting on context destruction

[kadircet]

• Drop LSP latency test

[kadircet]

as discussed offline there is no real use case for this now and it is unclear whether there will be.

dropping it, until we've got a use case.

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

 //===-- ClangdLSPServerTests.cpp ------------------------------------------===//
 //
 // 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 "Annotations.h"
 #include "ClangdLSPServer.h"
 #include "CodeComplete.h"
 #include "LSPClient.h"
 #include "Protocol.h"
 #include "TestFS.h"
 #include "refactor/Rename.h"
+#include "support/Context.h"
 #include "support/Logger.h"
 #include "support/TestTracer.h"
+#include "support/Threading.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Testing/Support/SupportHelpers.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
TEST_F(LSPTest, DiagnosticsHeaderSaved) {
   Client.notify(
       "textDocument/didSave",
       llvm::json::Object{{"textDocument", Client.documentID("foo.h")}});
   // Foo.cpp should be rebuilt with new diagnostics.
   EXPECT_THAT(Client.diagnostics("foo.cpp"),
               llvm::ValueIs(testing::ElementsAre(
                   DiagMessage("Use of undeclared identifier 'changed'"))));
 }
+
+TEST_F(LSPTest, RecordsLatencies) {
+  trace::TestTracer Tracer;
+  auto &Client = start();
+  llvm::StringLiteral MethodName = "method_name";
+  EXPECT_THAT(Tracer.takeMetric("lsp_latency", MethodName), testing::SizeIs(0));
+  llvm::consumeError(Client.call(MethodName, {}).take().takeError());
+  Client.sync();
+  EXPECT_THAT(Tracer.takeMetric("lsp_latency", MethodName), testing::SizeIs(1));
+}
 } // namespace
 } // namespace clangd
 } // namespace clang

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

 #include "LSPClient.h"
 #include "gtest/gtest.h"
 #include <condition_variable>
 
 #include "Protocol.h"
 #include "TestFS.h"
 #include "Transport.h"
+#include "support/Context.h"
 #include "support/Threading.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 #include <queue>
 
 namespace clang {
 namespace clangd {
 
std::pair<llvm::json::Value, CallResult *> addCallSlot() {
     unsigned ID = CallResults.size();
     CallResults.emplace_back();
     return {ID, &CallResults.back()};
   }
 
   // A null action causes the transport to shut down.
   void enqueue(std::function<void(MessageHandler &)> Action) {
     std::lock_guard<std::mutex> Lock(Mu);
-    Actions.push(std::move(Action));
+    Requests.push({Context::current().clone(), std::move(Action)});
     CV.notify_all();
   }
 
   std::vector<llvm::json::Value> takeNotifications(llvm::StringRef Method) {
     std::vector<llvm::json::Value> Result;
     {
       std::lock_guard<std::mutex> Lock(Mu);
       std::swap(Result, Notifications[Method]);
void call(llvm::StringRef Method, llvm::json::Value Params,
             llvm::json::Value ID) override {
     logBody(Method, Params, /*Send=*/false);
     ADD_FAILURE() << "Unexpected server->client call " << Method;
   }
 
   llvm::Error loop(MessageHandler &H) override {
     std::unique_lock<std::mutex> Lock(Mu);
     while (true) {
-      CV.wait(Lock, [&] { return !Actions.empty(); });
-      if (!Actions.front()) // Stop!
+      CV.wait(Lock, [&] { return !Requests.empty(); });
+      if (!Requests.front().Action) // Stop!
         return llvm::Error::success();
-      auto Action = std::move(Actions.front());
-      Actions.pop();
+      auto Req = std::move(Requests.front());
+      Requests.pop();
       Lock.unlock();
-      Action(H);
+      {
+        WithContext Ctx(std::move(Req.Ctx));
+        Req.Action(H);
+      }
       Lock.lock();
     }
   }
 
+  struct Request {
+    Context Ctx;
+    std::function<void(Transport::MessageHandler &)> Action;
+  };
+
   std::mutex Mu;
   std::deque<CallResult> CallResults;
-  std::queue<std::function<void(Transport::MessageHandler &)>> Actions;
+  std::queue<Request> Requests;
   std::condition_variable CV;
   llvm::StringMap<std::vector<llvm::json::Value>> Notifications;
 };
 
 LSPClient::LSPClient() : T(std::make_unique<TransportImpl>()) {}
 LSPClient::~LSPClient() = default;
 
 LSPClient::CallResult &LSPClient::call(llvm::StringRef Method,