sammccall created this revision.
sammccall added a reviewer: ilya-biryukov.
Herald added subscribers: cfe-commits, ioeric, jkorous-apple, klimek.

This is mostly to show what it would look like. I think it's pretty nice!

The lit tests pass, but the unit tests don't build because I ripped out some
std::future APIs and just return void, so the tests have no way to sync.
These probably need to be converted to callbacks first (callback APIs play
nicely with implicit contexts).


Repository:
  rCTE Clang Tools Extra

https://reviews.llvm.org/D42517

Files:
  clangd/ClangdLSPServer.cpp
  clangd/ClangdLSPServer.h
  clangd/ClangdServer.cpp
  clangd/ClangdServer.h
  clangd/ClangdUnit.cpp
  clangd/ClangdUnit.h
  clangd/CodeComplete.cpp
  clangd/CodeComplete.h
  clangd/Context.cpp
  clangd/Context.h
  clangd/Function.h
  clangd/GlobalCompilationDatabase.cpp
  clangd/JSONRPCDispatcher.cpp
  clangd/JSONRPCDispatcher.h
  clangd/Logger.cpp
  clangd/Logger.h
  clangd/ProtocolHandlers.cpp
  clangd/ProtocolHandlers.h
  clangd/Trace.cpp
  clangd/Trace.h
  clangd/XRefs.cpp
  clangd/XRefs.h
  clangd/index/FileIndex.cpp
  clangd/index/FileIndex.h
  clangd/index/Index.h
  clangd/index/MemIndex.cpp
  clangd/index/MemIndex.h
  clangd/index/Merge.cpp
  unittests/clangd/TraceTests.cpp

Index: unittests/clangd/TraceTests.cpp
===================================================================
--- unittests/clangd/TraceTests.cpp
+++ unittests/clangd/TraceTests.cpp
@@ -78,8 +78,8 @@
     auto JSONTracer = trace::createJSONTracer(OS);
     trace::Session Session(*JSONTracer);
     {
-      trace::Span S(Context::empty(), "A");
-      trace::log(Context::empty(), "B");
+      Context Ctx = trace::span(Context::empty(), "A");
+      trace::log(Ctx, "B");
     }
   }
 
Index: clangd/index/Merge.cpp
===================================================================
--- clangd/index/Merge.cpp
+++ clangd/index/Merge.cpp
@@ -24,7 +24,7 @@
    //          - find the generating file from each Symbol which is Static-only
    //          - ask Dynamic if it has that file (needs new SymbolIndex method)
    //          - if so, drop the Symbol.
-   bool fuzzyFind(const Context &Ctx, const FuzzyFindRequest &Req,
+   bool fuzzyFind(const FuzzyFindRequest &Req,
                   function_ref<void(const Symbol &)> Callback) const override {
      // We can't step through both sources in parallel. So:
      //  1) query all dynamic symbols, slurping results into a slab
@@ -34,13 +34,12 @@
      //  3) now yield all the dynamic symbols we haven't processed.
      bool More = false; // We'll be incomplete if either source was.
      SymbolSlab::Builder DynB;
-     More |=
-         Dynamic->fuzzyFind(Ctx, Req, [&](const Symbol &S) { DynB.insert(S); });
+     More |= Dynamic->fuzzyFind(Req, [&](const Symbol &S) { DynB.insert(S); });
      SymbolSlab Dyn = std::move(DynB).build();
 
      DenseSet<SymbolID> SeenDynamicSymbols;
      Symbol::Details Scratch;
-     More |= Static->fuzzyFind(Ctx, Req, [&](const Symbol &S) {
+     More |= Static->fuzzyFind(Req, [&](const Symbol &S) {
        auto DynS = Dyn.find(S.ID);
        if (DynS == Dyn.end())
          return Callback(S);
Index: clangd/index/MemIndex.h
===================================================================
--- clangd/index/MemIndex.h
+++ clangd/index/MemIndex.h
@@ -28,7 +28,7 @@
   static std::unique_ptr<SymbolIndex> build(SymbolSlab Slab);
 
   bool
-  fuzzyFind(const Context &Ctx, const FuzzyFindRequest &Req,
+  fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const override;
 
 private:
Index: clangd/index/MemIndex.cpp
===================================================================
--- clangd/index/MemIndex.cpp
+++ clangd/index/MemIndex.cpp
@@ -29,7 +29,7 @@
 }
 
 bool MemIndex::fuzzyFind(
-    const Context &Ctx, const FuzzyFindRequest &Req,
+    const FuzzyFindRequest &Req,
     llvm::function_ref<void(const Symbol &)> Callback) const {
   assert(!StringRef(Req.Query).contains("::") &&
          "There must be no :: in query.");
Index: clangd/index/Index.h
===================================================================
--- clangd/index/Index.h
+++ clangd/index/Index.h
@@ -10,7 +10,6 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_INDEX_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_INDEX_H
 
-#include "../Context.h"
 #include "clang/Index/IndexSymbol.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
@@ -244,7 +243,7 @@
   /// Returns true if the result list is complete, false if it was truncated due
   /// to MaxCandidateCount
   virtual bool
-  fuzzyFind(const Context &Ctx, const FuzzyFindRequest &Req,
+  fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const = 0;
 
   // FIXME: add interfaces for more index use cases:
Index: clangd/index/FileIndex.h
===================================================================
--- clangd/index/FileIndex.h
+++ clangd/index/FileIndex.h
@@ -17,7 +17,6 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H
 
 #include "../ClangdUnit.h"
-#include "../Context.h"
 #include "Index.h"
 #include "MemIndex.h"
 
@@ -58,10 +57,10 @@
 public:
   /// \brief Update symbols in \p Path with symbols in \p AST. If \p AST is
   /// nullptr, this removes all symbols in the file
-  void update(const Context &Ctx, PathRef Path, ParsedAST *AST);
+  void update(PathRef Path, ParsedAST *AST);
 
   bool
-  fuzzyFind(const Context &Ctx, const FuzzyFindRequest &Req,
+  fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const override;
 
 private:
Index: clangd/index/FileIndex.cpp
===================================================================
--- clangd/index/FileIndex.cpp
+++ clangd/index/FileIndex.cpp
@@ -69,7 +69,7 @@
   return {std::move(Snap), Pointers};
 }
 
-void FileIndex::update(const Context &Ctx, PathRef Path, ParsedAST *AST) {
+void FileIndex::update(PathRef Path, ParsedAST *AST) {
   if (!AST) {
     FSymbols.update(Path, nullptr);
   } else {
@@ -82,9 +82,9 @@
 }
 
 bool FileIndex::fuzzyFind(
-    const Context &Ctx, const FuzzyFindRequest &Req,
+    const FuzzyFindRequest &Req,
     llvm::function_ref<void(const Symbol &)> Callback) const {
-  return Index.fuzzyFind(Ctx, Req, Callback);
+  return Index.fuzzyFind(Req, Callback);
 }
 
 } // namespace clangd
Index: clangd/XRefs.h
===================================================================
--- clangd/XRefs.h
+++ clangd/XRefs.h
@@ -14,20 +14,18 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_XREFS_H
 
 #include "ClangdUnit.h"
-#include "Context.h"
 #include "Protocol.h"
 #include <vector>
 
 namespace clang {
 namespace clangd {
 
 /// Get definition of symbol at a specified \p Pos.
-std::vector<Location> findDefinitions(const Context &Ctx, ParsedAST &AST,
-                                      Position Pos);
+std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos);
 
 /// Returns highlights for all usages of a symbol at \p Pos.
-std::vector<DocumentHighlight>
-findDocumentHighlights(const Context &Ctx, ParsedAST &AST, Position Pos);
+std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
+                                                      Position Pos);
 
 } // namespace clangd
 } // namespace clang
Index: clangd/XRefs.cpp
===================================================================
--- clangd/XRefs.cpp
+++ clangd/XRefs.cpp
@@ -146,8 +146,7 @@
 
 } // namespace
 
-std::vector<Location> findDefinitions(const Context &Ctx, ParsedAST &AST,
-                                      Position Pos) {
+std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos) {
   const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
   const FileEntry *FE = SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
   if (!FE)
@@ -258,8 +257,8 @@
 
 } // namespace
 
-std::vector<DocumentHighlight>
-findDocumentHighlights(const Context &Ctx, ParsedAST &AST, Position Pos) {
+std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
+                                                      Position Pos) {
   const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
   const FileEntry *FE = SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
   if (!FE)
Index: clangd/Trace.h
===================================================================
--- clangd/Trace.h
+++ clangd/Trace.h
@@ -32,21 +32,17 @@
 /// Implmentations of this interface must be thread-safe.
 class EventTracer {
 public:
-  /// A callback executed when an event with duration ends. Args represent data
-  /// that was attached to the event via SPAN_ATTACH.
-  using EndEventCallback = UniqueFunction<void(json::obj &&Args)>;
-
   virtual ~EventTracer() = default;
 
-  /// Called when event that has a duration starts. The returned callback will
-  /// be executed when the event ends. \p Name is a descriptive name
-  /// of the event that was passed to Span constructor.
-  virtual EndEventCallback beginSpan(const Context &Ctx,
-                                     llvm::StringRef Name) = 0;
+  /// Called when event that has a duration starts. \p Name describes the event.
+  /// Returns a derived context that will be destroyed when the event ends.
+  /// Usually implementations will store an object in the returned context
+  /// whose destructor records the end of the event.
+  /// The args are spanArgs(), only complete when the event ends.
+  virtual Context beginSpan(llvm::StringRef Name) = 0;
 
   /// Called for instant events.
-  virtual void instant(const Context &Ctx, llvm::StringRef Name,
-                       json::obj &&Args) = 0;
+  virtual void instant(llvm::StringRef Name, json::obj &&Args) = 0;
 };
 
 /// Sets up a global EventTracer that consumes events produced by Span and
@@ -67,33 +63,30 @@
                                               bool Pretty = false);
 
 /// Records a single instant event, associated with the current thread.
-void log(const Context &Ctx, const llvm::Twine &Name);
+void log(const llvm::Twine &Name);
 
-/// Records an event whose duration is the lifetime of the Span object.
+/// Records an event with a duration.
+/// The event ends when the returned scoped object is destroyed.
+///
 /// This is the main public interface for producing tracing events.
 ///
 /// Arbitrary JSON metadata can be attached while this span is active:
-///   SPAN_ATTACH(MySpan, "Payload", SomeJSONExpr);
+///   SPAN_ATTACH(spanCtx, "Payload", SomeJSONExpr);
+///
 /// SomeJSONExpr is evaluated and copied only if actually needed.
-class Span {
-public:
-  Span(const Context &Ctx, llvm::StringRef Name);
-  ~Span();
-
-  /// Returns mutable span metadata if this span is interested.
-  /// Prefer to use SPAN_ATTACH rather than accessing this directly.
-  json::obj *args() { return Args.get(); }
+WithContext span(llvm::StringRef Name);
 
-private:
-  std::unique_ptr<json::obj> Args;
-  EventTracer::EndEventCallback Callback;
-};
+/// Returns mutable span metadata if this span is interested.
+/// Prefer to use SPAN_ATTACH rather than accessing this directly.
+json::obj *spanArgs();
 
-#define SPAN_ATTACH(S, Name, Expr)                                             \
+/// Attach a key-value pair to the current trace span.
+/// The context (or some ancestor) must be created by span().
+/// This macro is not safe for use on the same span from multiple threads.
+#define SPAN_ATTACH(Name, Expr)                                                \
   do {                                                                         \
-    if ((S).args() != nullptr) {                                               \
-      (*((S).args()))[(Name)] = (Expr);                                        \
-    }                                                                          \
+    if (auto *Args = ::clang::clangd::trace::spanArgs())                       \
+      (*Args)[Name] = Expr;                                                    \
   } while (0)
 
 } // namespace trace
Index: clangd/Trace.cpp
===================================================================
--- clangd/Trace.cpp
+++ clangd/Trace.cpp
@@ -7,6 +7,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "Context.h"
+#include "Function.h"
 #include "Trace.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/Chrono.h"
@@ -43,18 +45,15 @@
     Out.flush();
   }
 
-  EndEventCallback beginSpan(const Context &Ctx,
-                             llvm::StringRef Name) override {
+  Context beginSpan(llvm::StringRef Name) override {
     jsonEvent("B", json::obj{{"name", Name}});
-
-    // The callback that will run when event ends.
-    return [this](json::Expr &&Args) {
-      jsonEvent("E", json::obj{{"args", std::move(Args)}});
-    };
+    auto *Args = spanArgs();
+    return Context::current().derive(onScopeExit([this, Args] {
+      jsonEvent("E", json::obj{{"args", std::move(*Args)}});
+    }));
   }
 
-  void instant(const Context &Ctx, llvm::StringRef Name,
-               json::obj &&Args) override {
+  void instant(llvm::StringRef Name, json::obj &&Args) override {
     jsonEvent("i", json::obj{{"name", Name}, {"args", std::move(Args)}});
   }
 
@@ -119,29 +118,25 @@
   return llvm::make_unique<JSONTracer>(OS, Pretty);
 }
 
-void log(const Context &Ctx, const Twine &Message) {
+void log(const Twine &Message) {
   if (!T)
     return;
-  T->instant(Ctx, "Log", json::obj{{"Message", Message.str()}});
+  T->instant("Log", json::obj{{"Message", Message.str()}});
 }
 
-Span::Span(const Context &Ctx, llvm::StringRef Name) {
-  if (!T)
-    return;
-
-  Callback = T->beginSpan(Ctx, Name);
-  if (!Callback)
-    return;
+static Key<std::unique_ptr<json::obj>> kSpanArgs;
 
-  Args = llvm::make_unique<json::obj>();
+WithContext span(llvm::StringRef Name) {
+  if (!T)
+    return WithContext(Context::current().derive(kSpanArgs, nullptr));
+  return WithContext([Name] {
+    WithContextValue WithArgs(kSpanArgs, llvm::make_unique<json::obj>());
+    return T->beginSpan(Name);
+  }());
 }
 
-Span::~Span() {
-  if (!Callback)
-    return;
-
-  assert(Args && "Args must be non-null if Callback is defined");
-  Callback(std::move(*Args));
+json::obj *spanArgs() {
+  return Context::current().getExisting(kSpanArgs).get();
 }
 
 } // namespace trace
Index: clangd/ProtocolHandlers.h
===================================================================
--- clangd/ProtocolHandlers.h
+++ clangd/ProtocolHandlers.h
@@ -29,34 +29,28 @@
 // The interface implemented by ClangLSPServer to handle incoming requests.
 class ProtocolCallbacks {
 public:
-  // FIXME(ibiryukov): remove this typedef, inline its usages.
-  using Ctx = Context;
   virtual ~ProtocolCallbacks() = default;
 
-  virtual void onInitialize(Ctx C, InitializeParams &Params) = 0;
-  virtual void onShutdown(Ctx C, ShutdownParams &Params) = 0;
-  virtual void onExit(Ctx C, ExitParams &Params) = 0;
-  virtual void onDocumentDidOpen(Ctx C, DidOpenTextDocumentParams &Params) = 0;
-  virtual void onDocumentDidChange(Ctx C,
-                                   DidChangeTextDocumentParams &Params) = 0;
-  virtual void onDocumentDidClose(Ctx C,
-                                  DidCloseTextDocumentParams &Params) = 0;
-  virtual void onDocumentFormatting(Ctx C,
-                                    DocumentFormattingParams &Params) = 0;
+  virtual void onInitialize(InitializeParams &Params) = 0;
+  virtual void onShutdown(ShutdownParams &Params) = 0;
+  virtual void onExit(ExitParams &Params) = 0;
+  virtual void onDocumentDidOpen(DidOpenTextDocumentParams &Params) = 0;
+  virtual void onDocumentDidChange(DidChangeTextDocumentParams &Params) = 0;
+  virtual void onDocumentDidClose(DidCloseTextDocumentParams &Params) = 0;
+  virtual void onDocumentFormatting(DocumentFormattingParams &Params) = 0;
   virtual void
-  onDocumentOnTypeFormatting(Ctx C, DocumentOnTypeFormattingParams &Params) = 0;
+  onDocumentOnTypeFormatting(DocumentOnTypeFormattingParams &Params) = 0;
   virtual void
-  onDocumentRangeFormatting(Ctx C, DocumentRangeFormattingParams &Params) = 0;
-  virtual void onCodeAction(Ctx C, CodeActionParams &Params) = 0;
-  virtual void onCompletion(Ctx C, TextDocumentPositionParams &Params) = 0;
-  virtual void onSignatureHelp(Ctx C, TextDocumentPositionParams &Params) = 0;
-  virtual void onGoToDefinition(Ctx C, TextDocumentPositionParams &Params) = 0;
-  virtual void onSwitchSourceHeader(Ctx C, TextDocumentIdentifier &Params) = 0;
-  virtual void onFileEvent(Ctx C, DidChangeWatchedFilesParams &Params) = 0;
-  virtual void onCommand(Ctx C, ExecuteCommandParams &Params) = 0;
-  virtual void onRename(Ctx C, RenameParams &Parames) = 0;
-  virtual void onDocumentHighlight(Ctx C,
-                                   TextDocumentPositionParams &Params) = 0;
+  onDocumentRangeFormatting(DocumentRangeFormattingParams &Params) = 0;
+  virtual void onCodeAction(CodeActionParams &Params) = 0;
+  virtual void onCompletion(TextDocumentPositionParams &Params) = 0;
+  virtual void onSignatureHelp(TextDocumentPositionParams &Params) = 0;
+  virtual void onGoToDefinition(TextDocumentPositionParams &Params) = 0;
+  virtual void onSwitchSourceHeader(TextDocumentIdentifier &Params) = 0;
+  virtual void onFileEvent(DidChangeWatchedFilesParams &Params) = 0;
+  virtual void onCommand(ExecuteCommandParams &Params) = 0;
+  virtual void onRename(RenameParams &Parames) = 0;
+  virtual void onDocumentHighlight(TextDocumentPositionParams &Params) = 0;
 };
 
 void registerCallbackHandlers(JSONRPCDispatcher &Dispatcher, JSONOutput &Out,
Index: clangd/ProtocolHandlers.cpp
===================================================================
--- clangd/ProtocolHandlers.cpp
+++ clangd/ProtocolHandlers.cpp
@@ -24,17 +24,16 @@
 // FooParams should have a fromJSON function.
 struct HandlerRegisterer {
   template <typename Param>
-  void operator()(StringRef Method,
-                  void (ProtocolCallbacks::*Handler)(Context, Param)) {
+  void operator()(StringRef Method, void (ProtocolCallbacks::*Handler)(Param)) {
     // Capture pointers by value, as the lambda will outlive this object.
     auto *Callbacks = this->Callbacks;
     Dispatcher.registerHandler(
-        Method, [=](Context C, const json::Expr &RawParams) {
+        Method, [=](const json::Expr &RawParams) {
           typename std::remove_reference<Param>::type P;
           if (fromJSON(RawParams, P)) {
-            (Callbacks->*Handler)(std::move(C), P);
+            (Callbacks->*Handler)(P);
           } else {
-            log(C, "Failed to decode " + Method + " request.");
+            log("Failed to decode " + Method + " request.");
           }
         });
   }
Index: clangd/Logger.h
===================================================================
--- clangd/Logger.h
+++ clangd/Logger.h
@@ -10,24 +10,23 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_LOGGER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_LOGGER_H
 
-#include "Context.h"
 #include "llvm/ADT/Twine.h"
 
 namespace clang {
 namespace clangd {
 
 /// Main logging function.
 /// Logs messages to a global logger, which can be set up by LoggingSesssion.
 /// If no logger is registered, writes to llvm::errs().
-void log(const Context &Ctx, const llvm::Twine &Message);
+void log(const llvm::Twine &Message);
 
 /// Interface to allow custom logging in clangd.
 class Logger {
 public:
   virtual ~Logger() = default;
 
   /// Implementations of this method must be thread-safe.
-  virtual void log(const Context &Ctx, const llvm::Twine &Message) = 0;
+  virtual void log(const llvm::Twine &Message) = 0;
 };
 
 /// Only one LoggingSession can be active at a time.
Index: clangd/Logger.cpp
===================================================================
--- clangd/Logger.cpp
+++ clangd/Logger.cpp
@@ -25,9 +25,9 @@
 
 LoggingSession::~LoggingSession() { L = nullptr; }
 
-void log(const Context &Ctx, const llvm::Twine &Message) {
+void log(const llvm::Twine &Message) {
   if (L)
-    L->log(Ctx, Message);
+    L->log(Message);
   else {
     static std::mutex Mu;
     std::lock_guard<std::mutex> Guard(Mu);
Index: clangd/JSONRPCDispatcher.h
===================================================================
--- clangd/JSONRPCDispatcher.h
+++ clangd/JSONRPCDispatcher.h
@@ -10,7 +10,6 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_JSONRPCDISPATCHER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_JSONRPCDISPATCHER_H
 
-#include "Context.h"
 #include "JSONExpr.h"
 #include "Logger.h"
 #include "Protocol.h"
@@ -38,7 +37,7 @@
   void writeMessage(const json::Expr &Result);
 
   /// Write a line to the logging stream.
-  void log(const Context &Ctx, const Twine &Message) override;
+  void log(const Twine &Message) override;
 
   /// Mirror \p Message into InputMirror stream. Does nothing if InputMirror is
   /// null.
@@ -56,23 +55,22 @@
   std::mutex StreamMutex;
 };
 
-/// Sends a successful reply. \p Ctx must either be the Context accepted by
-/// JSONRPCDispatcher::Handler or be derived from it.
-void reply(const Context &Ctx, json::Expr &&Result);
-/// Sends an error response to the client, and logs it. \p Ctx must either be
-/// the Context accepted by JSONRPCDispatcher::Handler or be derived from it.
-void replyError(const Context &Ctx, ErrorCode code,
-                const llvm::StringRef &Message);
-/// Sends a request to the client. \p Ctx must either be the Context accepted by
-/// JSONRPCDispatcher::Handler or be derived from it.
-void call(const Context &Ctx, llvm::StringRef Method, json::Expr &&Params);
+/// Sends a successful reply.
+/// Current context must derive from JSONRPCDispatcher::Handler.
+void reply(json::Expr &&Result);
+/// Sends an error response to the client, and logs it.
+/// Current context must derive from JSONRPCDispatcher::Handler.
+void replyError(ErrorCode code, const llvm::StringRef &Message);
+/// Sends a request to the client.
+/// Current context must derive from JSONRPCDispatcher::Handler.
+void call(llvm::StringRef Method, json::Expr &&Params);
 
 /// Main JSONRPC entry point. This parses the JSONRPC "header" and calls the
 /// registered Handler for the method received.
 class JSONRPCDispatcher {
 public:
   // A handler responds to requests for a particular method name.
-  using Handler = std::function<void(Context, const json::Expr &)>;
+  using Handler = std::function<void(const json::Expr &)>;
 
   /// Create a new JSONRPCDispatcher. UnknownHandler is called when an unknown
   /// method is received.
Index: clangd/JSONRPCDispatcher.cpp
===================================================================
--- clangd/JSONRPCDispatcher.cpp
+++ clangd/JSONRPCDispatcher.cpp
@@ -20,7 +20,7 @@
 using namespace clangd;
 
 namespace {
-static Key<std::unique_ptr<trace::Span>> RequestSpan;
+static Key<Context> RequestCtx;
 static Key<json::Expr> RequestID;
 static Key<JSONOutput *> RequestOut;
 } // namespace
@@ -44,8 +44,8 @@
   Outs.flush();
 }
 
-void JSONOutput::log(const Context &Ctx, const Twine &Message) {
-  trace::log(Ctx, Message);
+void JSONOutput::log(const Twine &Message) {
+  trace::log(Message);
   std::lock_guard<std::mutex> Guard(StreamMutex);
   Logs << Message << '\n';
   Logs.flush();
@@ -59,34 +59,41 @@
   InputMirror->flush();
 }
 
-void clangd::reply(const Context &Ctx, json::Expr &&Result) {
-  auto ID = Ctx.get(RequestID);
+template <typename Func>
+static void InRequestContext(Func &&F) {
+  if (auto *EnclosingRequest = Context::current().get(RequestCtx)) {
+    WithContext Swap(EnclosingRequest->clone());
+    F();
+  }
+}
+
+void clangd::reply(json::Expr &&Result) {
+  auto ID = Context::current().get(RequestID);
   if (!ID) {
-    log(Ctx, "Attempted to reply to a notification!");
+    log("Attempted to reply to a notification!");
     return;
   }
 
-  if (auto *Span = Ctx.get(RequestSpan))
-    SPAN_ATTACH(**Span, "Reply", Result);
+  InRequestContext([&] { SPAN_ATTACH("Reply", Result); });
 
-  Ctx.getExisting(RequestOut)
+  Context::current().getExisting(RequestOut)
       ->writeMessage(json::obj{
           {"jsonrpc", "2.0"},
           {"id", *ID},
           {"result", std::move(Result)},
       });
 }
 
-void clangd::replyError(const Context &Ctx, ErrorCode code,
-                        const llvm::StringRef &Message) {
-  log(Ctx, "Error " + Twine(static_cast<int>(code)) + ": " + Message);
-  if (auto *Span = Ctx.get(RequestSpan))
-    SPAN_ATTACH(**Span, "Error",
-                (json::obj{{"code", static_cast<int>(code)},
-                           {"message", Message.str()}}));
+void clangd::replyError(ErrorCode code, const llvm::StringRef &Message) {
+  log("Error " + Twine(static_cast<int>(code)) + ": " + Message);
+  InRequestContext([&] {
+    SPAN_ATTACH("Error", (json::obj{{"code", static_cast<int>(code)},
+                                    {"message", Message.str()}}));
+  });
 
-  if (auto ID = Ctx.get(RequestID)) {
-    Ctx.getExisting(RequestOut)
+  if (auto ID = Context::current().get(RequestID)) {
+    Context::current()
+        .getExisting(RequestOut)
         ->writeMessage(json::obj{
             {"jsonrpc", "2.0"},
             {"id", *ID},
@@ -96,13 +103,15 @@
   }
 }
 
-void clangd::call(const Context &Ctx, StringRef Method, json::Expr &&Params) {
+void clangd::call(StringRef Method, json::Expr &&Params) {
   // FIXME: Generate/Increment IDs for every request so that we can get proper
   // replies once we need to.
-  if (auto *Span = Ctx.get(RequestSpan))
-    SPAN_ATTACH(**Span, "Call",
+  InRequestContext([&] {
+    SPAN_ATTACH("Call",
                 (json::obj{{"method", Method.str()}, {"params", Params}}));
-  Ctx.getExisting(RequestOut)
+  });
+  Context::current()
+      .getExisting(RequestOut)
       ->writeMessage(json::obj{
           {"jsonrpc", "2.0"},
           {"id", 1},
@@ -138,20 +147,20 @@
   auto &Handler = I != Handlers.end() ? I->second : UnknownHandler;
 
   // Create a Context that contains request information.
-  auto Ctx = Context::empty().derive(RequestOut, &Out);
+  WithContextValue WithRequestOut(RequestOut, &Out);
+  llvm::Optional<WithContextValue> WithID;
   if (ID)
-    Ctx = std::move(Ctx).derive(RequestID, *ID);
+    WithID.emplace(RequestID, *ID);
 
   // Create a tracing Span covering the whole request lifetime.
-  auto Tracer = llvm::make_unique<trace::Span>(Ctx, *Method);
+  auto Tracer = trace::span(*Method);
   if (ID)
-    SPAN_ATTACH(*Tracer, "ID", *ID);
-  SPAN_ATTACH(*Tracer, "Params", Params);
-
-  // Update Ctx to include Tracer.
-  Ctx = std::move(Ctx).derive(RequestSpan, std::move(Tracer));
+    SPAN_ATTACH("ID", *ID);
+  SPAN_ATTACH("Params", Params);
 
-  Handler(std::move(Ctx), std::move(Params));
+  // Stash the enclosing request context, so later calls can add metadata.
+  WithContextValue WithRequestContext(RequestCtx, Context::current().clone());
+  Handler(std::move(Params));
   return true;
 }
 
@@ -193,10 +202,9 @@
       // The end of headers is signified by an empty line.
       if (LineRef.consume_front("Content-Length: ")) {
         if (ContentLength != 0) {
-          log(Context::empty(),
-              "Warning: Duplicate Content-Length header received. "
+          log("Warning: Duplicate Content-Length header received. "
               "The previous value for this message (" +
-                  llvm::Twine(ContentLength) + ") was ignored.\n");
+              llvm::Twine(ContentLength) + ") was ignored.\n");
         }
 
         llvm::getAsUnsignedInteger(LineRef.trim(), 0, ContentLength);
@@ -215,8 +223,8 @@
     // and we don't want to crash downstream because of it.
     if (ContentLength > 1 << 30) { // 1024M
       In.ignore(ContentLength);
-      log(Context::empty(), "Skipped overly large message of " +
-                                Twine(ContentLength) + " bytes.\n");
+      log("Skipped overly large message of " + Twine(ContentLength) +
+          " bytes.\n");
       continue;
     }
 
@@ -230,35 +238,32 @@
         // If the stream is aborted before we read ContentLength bytes, In
         // will have eofbit and failbit set.
         if (!In) {
-          log(Context::empty(),
-              "Input was aborted. Read only " + llvm::Twine(In.gcount()) +
-                  " bytes of expected " + llvm::Twine(ContentLength) + ".\n");
+          log("Input was aborted. Read only " + llvm::Twine(In.gcount()) +
+              " bytes of expected " + llvm::Twine(ContentLength) + ".\n");
           break;
         }
 
         JSONRef = StringRef(JSON.data(), ContentLength);
       }
 
       if (auto Doc = json::parse(JSONRef)) {
         // Log the formatted message.
-        log(Context::empty(),
-            llvm::formatv(Out.Pretty ? "<-- {0:2}\n" : "<-- {0}\n", *Doc));
+        log(llvm::formatv(Out.Pretty ? "<-- {0:2}\n" : "<-- {0}\n", *Doc));
         // Finally, execute the action for this JSON message.
         if (!Dispatcher.call(*Doc, Out))
-          log(Context::empty(), "JSON dispatch failed!\n");
+          log("JSON dispatch failed!\n");
       } else {
         // Parse error. Log the raw message.
-        log(Context::empty(), "<-- " + JSONRef + "\n");
-        log(Context::empty(), llvm::Twine("JSON parse error: ") +
-                                  llvm::toString(Doc.takeError()) + "\n");
+        log("<-- " + JSONRef + "\n");
+        log(llvm::Twine("JSON parse error: ") +
+            llvm::toString(Doc.takeError()) + "\n");
       }
 
       // If we're done, exit the loop.
       if (IsDone)
         break;
     } else {
-      log(Context::empty(),
-          "Warning: Missing Content-Length header, or message has zero "
+      log("Warning: Missing Content-Length header, or message has zero "
           "length.\n");
     }
   }
Index: clangd/GlobalCompilationDatabase.cpp
===================================================================
--- clangd/GlobalCompilationDatabase.cpp
+++ clangd/GlobalCompilationDatabase.cpp
@@ -38,8 +38,7 @@
       return std::move(Candidates.front());
     }
   } else {
-    log(Context::empty(), // FIXME(ibiryukov): pass a proper Context here.
-        "Failed to find compilation database for " + Twine(File));
+    log("Failed to find compilation database for " + Twine(File));
   }
   return llvm::None;
 }
Index: clangd/Function.h
===================================================================
--- clangd/Function.h
+++ clangd/Function.h
@@ -144,7 +144,7 @@
   static_assert(std::is_same<typename std::decay<Func>::type, Func>::value,
                 "Func must be decayed");
 
-  ScopeExitGuard(Func F) : F(std::move(F)) {}
+  ScopeExitGuard(Func F) : F(llvm::make_unique<Func>(std::move(F))) {}
   ~ScopeExitGuard() {
     if (!F)
       return;
@@ -159,7 +159,7 @@
   ScopeExitGuard &operator=(ScopeExitGuard &&Other) = default;
 
 private:
-  llvm::Optional<Func> F;
+  std::unique_ptr<Func> F;
 };
 } // namespace detail
 
Index: clangd/Context.h
===================================================================
--- clangd/Context.h
+++ clangd/Context.h
@@ -16,15 +16,28 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CONTEXT_H_
 
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Compiler.h"
 #include <memory>
 #include <type_traits>
 
 namespace clang {
 namespace clangd {
 
-/// A key for a value of type \p Type, stored inside a context. Keys are
-/// non-movable and non-copyable. See documentation of the Context class for
-/// more details and usage examples.
+/// Values in a Context are indexed by typed keys.
+/// Key<T> serves two purposes:
+///   - it provides a lookup key for the context (each Key is unique),
+///   - it makes lookup type-safe: a Key<T> can only map to a T (or nothing).
+///
+/// Example:
+///    Key<int> RequestID;
+///    Key<int> Version;
+///
+///    Context Ctx = Context::empty().derive(RequestID, 10).derive(Version, 3);
+///    assert(*Ctx.get(RequestID) == 10);
+///    assert(*Ctx.get(Version) == 3);
+///
+/// Keys are typically used across multiple functions, so most of the time you
+/// would want to make them static class members or global variables.
 template <class Type> class Key {
 public:
   static_assert(!std::is_reference<Type>::value,
@@ -45,50 +58,23 @@
 /// Conceptually, a context is a heterogeneous map<Key<T>, T>. Each key has
 /// an associated value type, which allows the map to be typesafe.
 ///
+/// There is an "ambient" context for each thread, Context::current().
+/// Most functions should read from this, and use WithContextValue or
+/// WithContext to extend or replace the context within a block scope.
+/// Only code dealing with threads and extension points should need to use
+/// other Context objects.
+///
 /// You can't add data to an existing context, instead you create a new
 /// immutable context derived from it with extra data added. When you retrieve
 /// data, the context will walk up the parent chain until the key is found.
-///
-/// Contexts should be:
-///  - passed by reference when calling synchronous functions
-///  - passed by value (move) when calling asynchronous functions. The result
-///    callback of async operations will receive the context again.
-///  - cloned only when 'forking' an asynchronous computation that we don't wait
-///    for.
-///
-/// Copy operations for this class are deleted, use an explicit clone() method
-/// when you need a copy of the context instead.
-///
-/// To derive a child context use derive() function, e.g.
-///     Context ChildCtx = ParentCtx.derive(RequestIdKey, 123);
-///
-/// To create a new root context, derive() from empty Context.
-/// e.g.:
-///     Context Ctx = Context::empty().derive(RequestIdKey, 123);
-///
-/// Values in the context are indexed by typed keys (instances of Key<T> class).
-/// Key<T> serves two purposes:
-///   - it provides a lookup key for the context (each instance of a key is
-///   unique),
-///   - it keeps the type information about the value stored in the context map
-///   in the template arguments.
-/// This provides a type-safe interface to store and access values of multiple
-/// types inside a single context.
-/// For example,
-///    Key<int> RequestID;
-///    Key<int> Version;
-///
-///    Context Ctx = Context::empty().derive(RequestID, 10).derive(Version, 3);
-///    assert(*Ctx.get(RequestID) == 10);
-///    assert(*Ctx.get(Version) == 3);
-///
-/// Keys are typically used across multiple functions, so most of the time you
-/// would want to make them static class members or global variables.
 class Context {
 public:
-  /// Returns an empty context that contains no data. Useful for calling
-  /// functions that require a context when no explicit context is available.
+  /// Returns an empty root context that contains no data.
   static Context empty();
+  /// Returns the context for the current thread, creating it if needed.
+  /// Assigning to current() will permanently change the thread-local context.
+  /// See WithContext and WithContextValue for how to do this safely.
+  static Context &current();
 
 private:
   struct Data;
@@ -100,7 +86,8 @@
   /// (arguments of std::future<> must be default-construcitble in MSVC).
   Context() = default;
 
-  /// Move-only.
+  /// Copy operations for this class are deleted, use an explicit clone() method
+  /// when you need a copy of the context instead.
   Context(Context const &) = delete;
   Context &operator=(const Context &) = delete;
 
@@ -128,9 +115,8 @@
   }
 
   /// Derives a child context
-  /// It is safe to move or destroy a parent context after calling derive() from
-  /// it. The child context will continue to have access to the data stored in
-  /// the parent context.
+  /// It is safe to move or destroy a parent context after calling derive().
+  /// The child will keep its parent alive, and its data remains accessible.
   template <class Type>
   Context derive(const Key<Type> &Key,
                  typename std::decay<Type>::type Value) const & {
@@ -150,6 +136,18 @@
             std::move(Value))}));
   }
 
+  /// Derives a child context, using an anonymous key.
+  /// Intended for objects stored only for their destructor's side-effect.
+  template <class Type> Context derive(Type &&Value) const & {
+    static Key<typename std::decay<Type>::type> Private;
+    return derive(Private, std::forward<Type>(Value));
+  }
+
+  template <class Type> Context derive(Type &&Value) const && {
+    static Key<typename std::decay<Type>::type> Private;
+    return derive(Private, std::forward<Type>(Value));
+  }
+
   /// Clone this context object.
   Context clone() const;
 
@@ -183,7 +181,46 @@
   };
 
   std::shared_ptr<const Data> DataPtr;
-}; // namespace clangd
+};
+
+/// WithContext replaces Context::current() with a provided scope.
+/// When the WithContext is destroyed, the original scope is restored.
+/// For extending the current context with new value, prefer WithContextValue.
+class LLVM_NODISCARD WithContext {
+public:
+  WithContext(Context C)
+      : Restore(llvm::make_unique<Context>(std::move(Context::current()))) {
+    Context::current() = std::move(C);
+  }
+  ~WithContext() {
+    if (Restore)
+      Context::current() = std::move(*Restore);
+  }
+  WithContext(const WithContext&) = delete;
+  WithContext &operator=(const WithContext&) = delete;
+  WithContext(WithContext &&) = default; // Allow return-by-value.
+  WithContext &operator=(WithContext&&) = delete;
+
+private:
+  std::unique_ptr<Context> Restore; // unique_ptr for move semantics.
+};
+
+/// WithContextValue extends Context::current() with a single value.
+/// When the WithContextValue is destroyed, the original scope is restored.
+class LLVM_NODISCARD WithContextValue {
+public:
+  template <typename T>
+  WithContextValue(const Key<T> &K, typename std::decay<T>::type V)
+      : Restore(Context::current().derive(K, std::move(V))) {}
+
+  // Anonymous values can be used for the destructor side-effect.
+  template <typename T>
+  WithContextValue(T &&V)
+      : Restore(Context::current().derive(std::forward<T>(V))) {}
+
+private:
+  WithContext Restore;
+};
 
 } // namespace clangd
 } // namespace clang
Index: clangd/Context.cpp
===================================================================
--- clangd/Context.cpp
+++ clangd/Context.cpp
@@ -20,5 +20,10 @@
 
 Context Context::clone() const { return Context(DataPtr); }
 
+Context &Context::current() {
+  static thread_local Context C = empty();
+  return C;
+}
+
 } // namespace clangd
 } // namespace clang
Index: clangd/CodeComplete.h
===================================================================
--- clangd/CodeComplete.h
+++ clangd/CodeComplete.h
@@ -15,7 +15,6 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CODECOMPLETE_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CODECOMPLETE_H
 
-#include "Context.h"
 #include "Logger.h"
 #include "Path.h"
 #include "Protocol.h"
@@ -67,16 +66,16 @@
 };
 
 /// Get code completions at a specified \p Pos in \p FileName.
-CompletionList codeComplete(const Context &Ctx, PathRef FileName,
+CompletionList codeComplete(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
                             StringRef Contents, Position Pos,
                             IntrusiveRefCntPtr<vfs::FileSystem> VFS,
                             std::shared_ptr<PCHContainerOperations> PCHs,
                             CodeCompleteOptions Opts);
 
 /// Get signature help at a specified \p Pos in \p FileName.
-SignatureHelp signatureHelp(const Context &Ctx, PathRef FileName,
+SignatureHelp signatureHelp(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
                             StringRef Contents, Position Pos,
Index: clangd/CodeComplete.cpp
===================================================================
--- clangd/CodeComplete.cpp
+++ clangd/CodeComplete.cpp
@@ -636,8 +636,7 @@
 
 // Invokes Sema code completion on a file.
 // Callback will be invoked once completion is done, but before cleaning up.
-bool semaCodeComplete(const Context &Ctx,
-                      std::unique_ptr<CodeCompleteConsumer> Consumer,
+bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
                       const clang::CodeCompleteOptions &Options,
                       const SemaCompleteInput &Input,
                       llvm::function_ref<void()> Callback = nullptr) {
@@ -689,13 +688,12 @@
 
   SyntaxOnlyAction Action;
   if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0])) {
-    log(Ctx, "BeginSourceFile() failed when running codeComplete for " +
-                 Input.FileName);
+    log("BeginSourceFile() failed when running codeComplete for " +
+        Input.FileName);
     return false;
   }
   if (!Action.Execute()) {
-    log(Ctx,
-        "Execute() failed when running codeComplete for " + Input.FileName);
+    log("Execute() failed when running codeComplete for " + Input.FileName);
     return false;
   }
 
@@ -797,7 +795,6 @@
 //     This score is combined with the result quality score for the final score.
 //   - TopN determines the results with the best score.
 class CodeCompleteFlow {
-  const Context &Ctx;
   const CodeCompleteOptions &Opts;
   // Sema takes ownership of Recorder. Recorder is valid until Sema cleanup.
   std::unique_ptr<CompletionRecorder> RecorderOwner;
@@ -808,25 +805,24 @@
 
 public:
   // A CodeCompleteFlow object is only useful for calling run() exactly once.
-  CodeCompleteFlow(const Context &Ctx, const CodeCompleteOptions &Opts)
-      : Ctx(Ctx), Opts(Opts), RecorderOwner(new CompletionRecorder(Opts)),
+  CodeCompleteFlow(const CodeCompleteOptions &Opts)
+      : Opts(Opts), RecorderOwner(new CompletionRecorder(Opts)),
         Recorder(*RecorderOwner) {}
 
   CompletionList run(const SemaCompleteInput &SemaCCInput) && {
     // We run Sema code completion first. It builds an AST and calculates:
     //   - completion results based on the AST. These are saved for merging.
     //   - partial identifier and context. We need these for the index query.
     CompletionList Output;
-    semaCodeComplete(Ctx, std::move(RecorderOwner), Opts.getClangCompleteOpts(),
+    semaCodeComplete(std::move(RecorderOwner), Opts.getClangCompleteOpts(),
                      SemaCCInput, [&] {
                        if (Recorder.CCSema)
                          Output = runWithSema();
                        else
-                         log(Ctx, "Code complete: no Sema callback, 0 results");
+                         log("Code complete: no Sema callback, 0 results");
                      });
 
-    log(Ctx,
-        llvm::formatv("Code complete: {0} results from Sema, {1} from Index, "
+    log(llvm::formatv("Code complete: {0} results from Sema, {1} from Index, "
                       "{2} matched, {3} returned{4}.",
                       NSema, NIndex, NBoth, Output.items.size(),
                       Output.isIncomplete ? " (incomplete)" : ""));
@@ -867,12 +863,12 @@
     Req.Query = Filter->pattern();
     Req.Scopes =
         getQueryScopes(Recorder.CCContext, Recorder.CCSema->getSourceManager());
-    log(Ctx, llvm::formatv(
-                 "Code complete: fuzzyFind(\"{0}\", Scopes: [{1}]", Req.Query,
-                 llvm::join(Req.Scopes.begin(), Req.Scopes.end(), ",")));
+    log(llvm::formatv("Code complete: fuzzyFind(\"{0}\", Scopes: [{1}]",
+                      Req.Query,
+                      llvm::join(Req.Scopes.begin(), Req.Scopes.end(), ",")));
     // Run the query against the index.
     Incomplete |= !Opts.Index->fuzzyFind(
-        Ctx, Req, [&](const Symbol &Sym) { ResultsBuilder.insert(Sym); });
+        Req, [&](const Symbol &Sym) { ResultsBuilder.insert(Sym); });
     return std::move(ResultsBuilder).build();
   }
 
@@ -941,18 +937,18 @@
   }
 };
 
-CompletionList codeComplete(const Context &Ctx, PathRef FileName,
+CompletionList codeComplete(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
                             StringRef Contents, Position Pos,
                             IntrusiveRefCntPtr<vfs::FileSystem> VFS,
                             std::shared_ptr<PCHContainerOperations> PCHs,
                             CodeCompleteOptions Opts) {
-  return CodeCompleteFlow(Ctx, Opts).run(
+  return CodeCompleteFlow(Opts).run(
       {FileName, Command, Preamble, Contents, Pos, VFS, PCHs});
 }
 
-SignatureHelp signatureHelp(const Context &Ctx, PathRef FileName,
+SignatureHelp signatureHelp(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
                             StringRef Contents, Position Pos,
@@ -964,10 +960,10 @@
   Options.IncludeMacros = false;
   Options.IncludeCodePatterns = false;
   Options.IncludeBriefComments = true;
-  semaCodeComplete(
-      Ctx, llvm::make_unique<SignatureHelpCollector>(Options, Result), Options,
-      {FileName, Command, Preamble, Contents, Pos, std::move(VFS),
-       std::move(PCHs)});
+  semaCodeComplete(llvm::make_unique<SignatureHelpCollector>(Options, Result),
+                   Options,
+                   {FileName, Command, Preamble, Contents, Pos, std::move(VFS),
+                    std::move(PCHs)});
   return Result;
 }
 
Index: clangd/ClangdUnit.h
===================================================================
--- clangd/ClangdUnit.h
+++ clangd/ClangdUnit.h
@@ -10,7 +10,6 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNIT_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNIT_H
 
-#include "Context.h"
 #include "Function.h"
 #include "Path.h"
 #include "Protocol.h"
@@ -71,7 +70,7 @@
   /// Attempts to run Clang and store parsed AST. If \p Preamble is non-null
   /// it is reused during parsing.
   static llvm::Optional<ParsedAST>
-  Build(const Context &Ctx, std::unique_ptr<clang::CompilerInvocation> CI,
+  Build(std::unique_ptr<clang::CompilerInvocation> CI,
         std::shared_ptr<const PreambleData> Preamble,
         std::unique_ptr<llvm::MemoryBuffer> Buffer,
         std::shared_ptr<PCHContainerOperations> PCHs,
@@ -144,8 +143,7 @@
   mutable llvm::Optional<ParsedAST> AST;
 };
 
-using ASTParsedCallback =
-    std::function<void(const Context &Ctx, PathRef Path, ParsedAST *)>;
+using ASTParsedCallback = std::function<void(PathRef Path, ParsedAST *)>;
 
 /// Manages resources, required by clangd. Allows to rebuild file with new
 /// contents, and provides AST and Preamble for it.
@@ -182,8 +180,7 @@
   /// Returns a list of diagnostics or a llvm::None, if another rebuild was
   /// requested in parallel (effectively cancelling this rebuild) before
   /// diagnostics were produced.
-  llvm::Optional<std::vector<DiagWithFixIts>> rebuild(const Context &Ctx,
-                                                      ParseInputs &&Inputs);
+  llvm::Optional<std::vector<DiagWithFixIts>> rebuild(ParseInputs &&Inputs);
 
   /// Schedule a rebuild and return a deferred computation that will finish the
   /// rebuild, that can be called on a different thread.
@@ -199,7 +196,7 @@
   /// The future to finish rebuild returns a list of diagnostics built during
   /// reparse, or None, if another deferRebuild was called before this
   /// rebuild was finished.
-  UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>(const Context &)>
+  UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>()>
   deferRebuild(ParseInputs &&Inputs);
 
   /// Returns a future to get the most fresh PreambleData for a file. The
Index: clangd/ClangdUnit.cpp
===================================================================
--- clangd/ClangdUnit.cpp
+++ clangd/ClangdUnit.cpp
@@ -225,8 +225,7 @@
 }
 
 llvm::Optional<ParsedAST>
-ParsedAST::Build(const Context &Ctx,
-                 std::unique_ptr<clang::CompilerInvocation> CI,
+ParsedAST::Build(std::unique_ptr<clang::CompilerInvocation> CI,
                  std::shared_ptr<const PreambleData> Preamble,
                  std::unique_ptr<llvm::MemoryBuffer> Buffer,
                  std::shared_ptr<PCHContainerOperations> PCHs,
@@ -248,12 +247,12 @@
   auto Action = llvm::make_unique<ClangdFrontendAction>();
   const FrontendInputFile &MainInput = Clang->getFrontendOpts().Inputs[0];
   if (!Action->BeginSourceFile(*Clang, MainInput)) {
-    log(Ctx, "BeginSourceFile() failed when building AST for " +
+    log("BeginSourceFile() failed when building AST for " +
                  MainInput.getFile());
     return llvm::None;
   }
   if (!Action->Execute())
-    log(Ctx, "Execute() failed when building AST for " + MainInput.getFile());
+    log("Execute() failed when building AST for " + MainInput.getFile());
 
   // UnitDiagsConsumer is local, we can not store it in CompilerInstance that
   // has a longer lifetime.
@@ -372,8 +371,7 @@
     : FileName(FileName), StorePreamblesInMemory(StorePreamblesInMemory),
       RebuildCounter(0), RebuildInProgress(false), PCHs(std::move(PCHs)),
       ASTCallback(std::move(ASTCallback)) {
-  // FIXME(ibiryukov): we should pass a proper Context here.
-  log(Context::empty(), "Created CppFile for " + FileName);
+  log("Created CppFile for " + FileName);
 
   std::lock_guard<std::mutex> Lock(Mutex);
   LatestAvailablePreamble = nullptr;
@@ -425,11 +423,11 @@
 }
 
 llvm::Optional<std::vector<DiagWithFixIts>>
-CppFile::rebuild(const Context &Ctx, ParseInputs &&Inputs) {
-  return deferRebuild(std::move(Inputs))(Ctx);
+CppFile::rebuild(ParseInputs &&Inputs) {
+  return deferRebuild(std::move(Inputs))();
 }
 
-UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>(const Context &)>
+UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>()>
 CppFile::deferRebuild(ParseInputs &&Inputs) {
   std::shared_ptr<const PreambleData> OldPreamble;
   std::shared_ptr<PCHContainerOperations> PCHs;
@@ -467,13 +465,11 @@
   std::shared_ptr<CppFile> That = shared_from_this();
   auto FinishRebuild =
       [OldPreamble, RequestRebuildCounter, PCHs,
-       That](ParseInputs Inputs,
-             const Context &Ctx) mutable /* to allow changing OldPreamble. */
+       That](ParseInputs Inputs) mutable /* to allow changing OldPreamble. */
       -> llvm::Optional<std::vector<DiagWithFixIts>> {
-    log(Context::empty(),
-        "Rebuilding file " + That->FileName + " with command [" +
-            Inputs.CompileCommand.Directory + "] " +
-            llvm::join(Inputs.CompileCommand.CommandLine, " "));
+    log("Rebuilding file " + That->FileName + " with command [" +
+        Inputs.CompileCommand.Directory + "] " +
+        llvm::join(Inputs.CompileCommand.CommandLine, " "));
 
     // Only one execution of this method is possible at a time.
     // RebuildGuard will wait for any ongoing rebuilds to finish and will put us
@@ -516,17 +512,17 @@
       if (OldPreamble &&
           OldPreamble->Preamble.CanReuse(*CI, ContentsBuffer.get(), Bounds,
                                          Inputs.FS.get())) {
-        log(Ctx, "Reusing preamble for file " + Twine(That->FileName));
+        log("Reusing preamble for file " + Twine(That->FileName));
         return OldPreamble;
       }
-      log(Ctx, "Premble for file " + Twine(That->FileName) +
+      log("Premble for file " + Twine(That->FileName) +
                    " cannot be reused. Attempting to rebuild it.");
       // We won't need the OldPreamble anymore, release it so it can be
       // deleted (if there are no other references to it).
       OldPreamble.reset();
 
-      trace::Span Tracer(Ctx, "Preamble");
-      SPAN_ATTACH(Tracer, "File", That->FileName);
+      auto Tracer = trace::span("Preamble");
+      SPAN_ATTACH("File", That->FileName);
       std::vector<DiagWithFixIts> PreambleDiags;
       StoreDiagsConsumer PreambleDiagnosticsConsumer(/*ref*/ PreambleDiags);
       IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
@@ -550,16 +546,15 @@
       CI->getFrontendOpts().SkipFunctionBodies = false;
 
       if (BuiltPreamble) {
-        log(Ctx, "Built preamble of size " + Twine(BuiltPreamble->getSize()) +
-                     " for file " + Twine(That->FileName));
+        log("Built preamble of size " + Twine(BuiltPreamble->getSize()) +
+            " for file " + Twine(That->FileName));
 
         return std::make_shared<PreambleData>(
             std::move(*BuiltPreamble),
             SerializedDeclsCollector.takeTopLevelDeclIDs(),
             std::move(PreambleDiags));
       } else {
-        log(Ctx,
-            "Could not build a preamble for file " + Twine(That->FileName));
+        log("Could not build a preamble for file " + Twine(That->FileName));
         return nullptr;
       }
     };
@@ -587,17 +582,17 @@
     // Compute updated AST.
     llvm::Optional<ParsedAST> NewAST;
     {
-      trace::Span Tracer(Ctx, "Build");
-      SPAN_ATTACH(Tracer, "File", That->FileName);
-      NewAST = ParsedAST::Build(Ctx, std::move(CI), std::move(NewPreamble),
+      auto Tracer = trace::span("Build");
+      SPAN_ATTACH("File", That->FileName);
+      NewAST = ParsedAST::Build(std::move(CI), std::move(NewPreamble),
                                 std::move(ContentsBuffer), PCHs, Inputs.FS);
     }
 
     if (NewAST) {
       Diagnostics.insert(Diagnostics.end(), NewAST->getDiagnostics().begin(),
                          NewAST->getDiagnostics().end());
       if (That->ASTCallback)
-        That->ASTCallback(Ctx, That->FileName, NewAST.getPointer());
+        That->ASTCallback(That->FileName, NewAST.getPointer());
     } else {
       // Don't report even Preamble diagnostics if we coulnd't build AST.
       Diagnostics.clear();
Index: clangd/ClangdServer.h
===================================================================
--- clangd/ClangdServer.h
+++ clangd/ClangdServer.h
@@ -10,6 +10,7 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDSERVER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDSERVER_H
 
+#include "Context.h"
 #include "ClangdUnit.h"
 #include "ClangdUnitStore.h"
 #include "CodeComplete.h"
@@ -72,7 +73,7 @@
 
   /// Called by ClangdServer when \p Diagnostics for \p File are ready.
   virtual void
-  onDiagnosticsReady(const Context &Ctx, PathRef File,
+  onDiagnosticsReady(PathRef File,
                      Tagged<std::vector<DiagWithFixIts>> Diagnostics) = 0;
 };
 
@@ -125,7 +126,8 @@
 
     {
       std::lock_guard<std::mutex> Lock(Mutex);
-      RequestQueue.push_front(
+      RequestQueue.emplace_front(
+          Context::current().clone(),
           BindWithForward(std::forward<Func>(F), std::forward<Args>(As)...));
     }
     RequestCV.notify_one();
@@ -141,7 +143,8 @@
 
     {
       std::lock_guard<std::mutex> Lock(Mutex);
-      RequestQueue.push_back(
+      RequestQueue.emplace_back(
+          Context::current().clone(),
           BindWithForward(std::forward<Func>(F), std::forward<Args>(As)...));
     }
     RequestCV.notify_one();
@@ -158,7 +161,7 @@
   bool Done = false;
   /// A queue of requests. Elements of this vector are async computations (i.e.
   /// results of calling std::async(std::launch::deferred, ...)).
-  std::deque<UniqueFunction<void()>> RequestQueue;
+  std::deque<std::pair<Context, UniqueFunction<void()>>> RequestQueue;
   /// Condition variable to wake up worker threads.
   std::condition_variable RequestCV;
 };
@@ -220,26 +223,20 @@
   /// constructor will receive onDiagnosticsReady callback.
   /// \return A future that will become ready when the rebuild (including
   /// diagnostics) is finished.
-  std::future<Context> addDocument(Context Ctx, PathRef File,
-                                   StringRef Contents);
+  void addDocument(PathRef File, StringRef Contents);
   /// Remove \p File from list of tracked files, schedule a request to free
   /// resources associated with it.
   /// \return A future that will become ready when the file is removed and all
   /// associated resources are freed.
-  std::future<Context> removeDocument(Context Ctx, PathRef File);
+  void removeDocument(PathRef File);
   /// Force \p File to be reparsed using the latest contents.
   /// Will also check if CompileCommand, provided by GlobalCompilationDatabase
   /// for \p File has changed. If it has, will remove currently stored Preamble
   /// and AST and rebuild them from scratch.
-  std::future<Context> forceReparse(Context Ctx, PathRef File);
+  void forceReparse(PathRef File);
 
-  /// DEPRECATED. Please use a callback-based version, this API is deprecated
-  /// and will soon be removed.
-  ///
   /// Run code completion for \p File at \p Pos.
-  ///
-  /// Request is processed asynchronously. You can use the returned future to
-  /// wait for the results of the async request.
+  /// Request is processed asynchronously.
   ///
   /// If \p OverridenContents is not None, they will used only for code
   /// completion, i.e. no diagnostics update will be scheduled and a draft for
@@ -250,18 +247,12 @@
   /// This method should only be called for currently tracked files. However, it
   /// is safe to call removeDocument for \p File after this method returns, even
   /// while returned future is not yet ready.
-  std::future<std::pair<Context, Tagged<CompletionList>>>
-  codeComplete(Context Ctx, PathRef File, Position Pos,
-               const clangd::CodeCompleteOptions &Opts,
-               llvm::Optional<StringRef> OverridenContents = llvm::None,
-               IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS = nullptr);
-
   /// A version of `codeComplete` that runs \p Callback on the processing thread
   /// when codeComplete results become available.
   void
-  codeComplete(Context Ctx, PathRef File, Position Pos,
+  codeComplete(PathRef File, Position Pos,
                const clangd::CodeCompleteOptions &Opts,
-               UniqueFunction<void(Context, Tagged<CompletionList>)> Callback,
+               UniqueFunction<void(Tagged<CompletionList>)> Callback,
                llvm::Optional<StringRef> OverridenContents = llvm::None,
                IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS = nullptr);
 
@@ -273,21 +264,21 @@
   /// vfs::FileSystem used for signature help. This method should only be called
   /// for currently tracked files.
   llvm::Expected<Tagged<SignatureHelp>>
-  signatureHelp(const Context &Ctx, PathRef File, Position Pos,
+  signatureHelp(PathRef File, Position Pos,
                 llvm::Optional<StringRef> OverridenContents = llvm::None,
                 IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS = nullptr);
 
   /// Get definition of symbol at a specified \p Line and \p Column in \p File.
   llvm::Expected<Tagged<std::vector<Location>>>
-  findDefinitions(const Context &Ctx, PathRef File, Position Pos);
+  findDefinitions(PathRef File, Position Pos);
 
   /// Helper function that returns a path to the corresponding source file when
   /// given a header file and vice versa.
   llvm::Optional<Path> switchSourceHeader(PathRef Path);
 
   /// Get document highlights for a given position.
   llvm::Expected<Tagged<std::vector<DocumentHighlight>>>
-  findDocumentHighlights(const Context &Ctx, PathRef File, Position Pos);
+  findDocumentHighlights(PathRef File, Position Pos);
 
   /// Run formatting for \p Rng inside \p File with content \p Code.
   llvm::Expected<tooling::Replacements> formatRange(StringRef Code,
@@ -304,8 +295,7 @@
 
   /// Rename all occurrences of the symbol at the \p Pos in \p File to
   /// \p NewName.
-  Expected<std::vector<tooling::Replacement>> rename(const Context &Ctx,
-                                                     PathRef File, Position Pos,
+  Expected<std::vector<tooling::Replacement>> rename(PathRef File, Position Pos,
                                                      llvm::StringRef NewName);
 
   /// Gets current document contents for \p File. Returns None if \p File is not
@@ -328,14 +318,13 @@
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
-  std::future<Context>
-  scheduleReparseAndDiags(Context Ctx, PathRef File, VersionedDraft Contents,
+  void
+  scheduleReparseAndDiags(PathRef File, VersionedDraft Contents,
                           std::shared_ptr<CppFile> Resources,
                           Tagged<IntrusiveRefCntPtr<vfs::FileSystem>> TaggedFS,
                           bool AllowCachedCompileFlags);
 
-  std::future<Context>
-  scheduleCancelRebuild(Context Ctx, std::shared_ptr<CppFile> Resources);
+  void scheduleCancelRebuild(std::shared_ptr<CppFile> Resources);
 
   GlobalCompilationDatabase &CDB;
   DiagnosticsConsumer &DiagConsumer;
Index: clangd/ClangdServer.cpp
===================================================================
--- clangd/ClangdServer.cpp
+++ clangd/ClangdServer.cpp
@@ -101,6 +101,7 @@
       llvm::set_thread_name(llvm::formatv("scheduler/{0}", I));
       while (true) {
         UniqueFunction<void()> Request;
+        llvm::Optional<WithContext> Ctx;
 
         // Pick request from the queue
         {
@@ -117,10 +118,10 @@
           // ClangdScheduler have a way to prioritise their requests by putting
           // them to the either side of the queue (using either addToEnd or
           // addToFront).
-          Request = std::move(RequestQueue.front());
+          Request = std::move(RequestQueue.front().second);
+          Ctx.emplace(std::move(RequestQueue.front().first));
           RequestQueue.pop_front();
         } // unlock Mutex
-
         Request();
       }
     }));
@@ -156,10 +157,9 @@
       // parsed.
       // FIXME(ioeric): this can be slow and we may be able to index on less
       // critical paths.
-      Units(FileIdx
-                ? [this](const Context &Ctx, PathRef Path,
-                         ParsedAST *AST) { FileIdx->update(Ctx, Path, AST); }
-                : ASTParsedCallback()),
+      Units(FileIdx ? [this](PathRef Path,
+                             ParsedAST *AST) { FileIdx->update(Path, AST); }
+                    : ASTParsedCallback()),
       ResourceDir(ResourceDir ? ResourceDir->str() : getStandardResourceDir()),
       PCHs(std::make_shared<PCHContainerOperations>()),
       StorePreamblesInMemory(StorePreamblesInMemory),
@@ -182,66 +182,42 @@
     this->RootPath = NewRootPath;
 }
 
-std::future<Context> ClangdServer::addDocument(Context Ctx, PathRef File,
-                                               StringRef Contents) {
+void ClangdServer::addDocument(PathRef File, StringRef Contents) {
   DocVersion Version = DraftMgr.updateDraft(File, Contents);
 
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
   std::shared_ptr<CppFile> Resources =
       Units.getOrCreateFile(File, ResourceDir, StorePreamblesInMemory, PCHs);
-  return scheduleReparseAndDiags(std::move(Ctx), File,
-                                 VersionedDraft{Version, Contents.str()},
+  return scheduleReparseAndDiags(File, VersionedDraft{Version, Contents.str()},
                                  std::move(Resources), std::move(TaggedFS),
                                  /*AllowCachedCompileFlags=*/true);
 }
 
-std::future<Context> ClangdServer::removeDocument(Context Ctx, PathRef File) {
+void ClangdServer::removeDocument(PathRef File) {
   DraftMgr.removeDraft(File);
   std::shared_ptr<CppFile> Resources = Units.removeIfPresent(File);
-  return scheduleCancelRebuild(std::move(Ctx), std::move(Resources));
+  return scheduleCancelRebuild(std::move(Resources));
 }
 
-std::future<Context> ClangdServer::forceReparse(Context Ctx, PathRef File) {
+void ClangdServer::forceReparse(PathRef File) {
   auto FileContents = DraftMgr.getDraft(File);
   assert(FileContents.Draft &&
          "forceReparse() was called for non-added document");
 
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
   std::shared_ptr<CppFile> Resources =
       Units.getOrCreateFile(File, ResourceDir, StorePreamblesInMemory, PCHs);
-  return scheduleReparseAndDiags(std::move(Ctx), File, FileContents,
-                                 std::move(Resources), std::move(TaggedFS),
+  return scheduleReparseAndDiags(File, FileContents, std::move(Resources),
+                                 std::move(TaggedFS),
                                  /*AllowCachedCompileFlags=*/false);
 }
 
-std::future<std::pair<Context, Tagged<CompletionList>>>
-ClangdServer::codeComplete(Context Ctx, PathRef File, Position Pos,
-                           const clangd::CodeCompleteOptions &Opts,
-                           llvm::Optional<StringRef> OverridenContents,
-                           IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS) {
-  using ResultType = std::pair<Context, Tagged<CompletionList>>;
-
-  std::promise<ResultType> ResultPromise;
-
-  auto Callback = [](std::promise<ResultType> ResultPromise, Context Ctx,
-                     Tagged<CompletionList> Result) -> void {
-    ResultPromise.set_value({std::move(Ctx), std::move(Result)});
-  };
-
-  std::future<ResultType> ResultFuture = ResultPromise.get_future();
-  codeComplete(std::move(Ctx), File, Pos, Opts,
-               BindWithForward(Callback, std::move(ResultPromise)),
-               OverridenContents, UsedFS);
-  return ResultFuture;
-}
-
 void ClangdServer::codeComplete(
-    Context Ctx, PathRef File, Position Pos,
-    const clangd::CodeCompleteOptions &Opts,
-    UniqueFunction<void(Context, Tagged<CompletionList>)> Callback,
+    PathRef File, Position Pos, const clangd::CodeCompleteOptions &Opts,
+    UniqueFunction<void(Tagged<CompletionList>)> Callback,
     llvm::Optional<StringRef> OverridenContents,
     IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS) {
-  using CallbackType = UniqueFunction<void(Context, Tagged<CompletionList>)>;
+  using CallbackType = UniqueFunction<void(Tagged<CompletionList>)>;
 
   std::string Contents;
   if (OverridenContents) {
@@ -286,28 +262,26 @@
   auto Task =
       // 'mutable' to reassign Preamble variable.
       [FileStr, Preamble, Resources, Contents, Pos, CodeCompleteOpts, TaggedFS,
-       PCHs, CompileCommand](Context Ctx, CallbackType Callback) mutable {
+       PCHs, CompileCommand](CallbackType Callback) mutable {
         if (!Preamble) {
           // Maybe we built some preamble before processing this request.
           Preamble = Resources->getPossiblyStalePreamble();
         }
         // FIXME(ibiryukov): even if Preamble is non-null, we may want to check
         // both the old and the new version in case only one of them matches.
         CompletionList Result = clangd::codeComplete(
-            Ctx, FileStr, CompileCommand,
+            FileStr, CompileCommand,
             Preamble ? &Preamble->Preamble : nullptr, Contents, Pos,
             TaggedFS.Value, PCHs, CodeCompleteOpts);
 
-        Callback(std::move(Ctx),
-                 make_tagged(std::move(Result), std::move(TaggedFS.Tag)));
+        Callback(make_tagged(std::move(Result), std::move(TaggedFS.Tag)));
       };
 
-  WorkScheduler.addToFront(std::move(Task), std::move(Ctx),
-                           std::move(Callback));
+  WorkScheduler.addToFront(std::move(Task), std::move(Callback));
 }
 
 llvm::Expected<Tagged<SignatureHelp>>
-ClangdServer::signatureHelp(const Context &Ctx, PathRef File, Position Pos,
+ClangdServer::signatureHelp(PathRef File, Position Pos,
                             llvm::Optional<StringRef> OverridenContents,
                             IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS) {
   std::string DraftStorage;
@@ -337,7 +311,7 @@
 
   auto Preamble = Resources->getPossiblyStalePreamble();
   auto Result =
-      clangd::signatureHelp(Ctx, File, *Resources->getLastCommand(),
+      clangd::signatureHelp(File, *Resources->getLastCommand(),
                             Preamble ? &Preamble->Preamble : nullptr,
                             *OverridenContents, Pos, TaggedFS.Value, PCHs);
   return make_tagged(std::move(Result), TaggedFS.Tag);
@@ -370,8 +344,7 @@
 }
 
 Expected<std::vector<tooling::Replacement>>
-ClangdServer::rename(const Context &Ctx, PathRef File, Position Pos,
-                     llvm::StringRef NewName) {
+ClangdServer::rename(PathRef File, Position Pos, llvm::StringRef NewName) {
   std::shared_ptr<CppFile> Resources = Units.getFile(File);
   RefactoringResultCollector ResultCollector;
   Resources->getAST().get()->runUnderLock([&](ParsedAST *AST) {
@@ -443,7 +416,7 @@
 }
 
 llvm::Expected<Tagged<std::vector<Location>>>
-ClangdServer::findDefinitions(const Context &Ctx, PathRef File, Position Pos) {
+ClangdServer::findDefinitions(PathRef File, Position Pos) {
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
 
   std::shared_ptr<CppFile> Resources = Units.getFile(File);
@@ -453,10 +426,10 @@
         llvm::errc::invalid_argument);
 
   std::vector<Location> Result;
-  Resources->getAST().get()->runUnderLock([Pos, &Result, &Ctx](ParsedAST *AST) {
+  Resources->getAST().get()->runUnderLock([Pos, &Result](ParsedAST *AST) {
     if (!AST)
       return;
-    Result = clangd::findDefinitions(Ctx, *AST, Pos);
+    Result = clangd::findDefinitions(*AST, Pos);
   });
   return make_tagged(std::move(Result), TaggedFS.Tag);
 }
@@ -535,8 +508,7 @@
 }
 
 llvm::Expected<Tagged<std::vector<DocumentHighlight>>>
-ClangdServer::findDocumentHighlights(const Context &Ctx, PathRef File,
-                                     Position Pos) {
+ClangdServer::findDocumentHighlights(PathRef File, Position Pos) {
   auto FileContents = DraftMgr.getDraft(File);
   if (!FileContents.Draft)
     return llvm::make_error<llvm::StringError>(
@@ -553,24 +525,22 @@
 
   std::vector<DocumentHighlight> Result;
   llvm::Optional<llvm::Error> Err;
-  Resources->getAST().get()->runUnderLock([Pos, &Ctx, &Err,
-                                           &Result](ParsedAST *AST) {
+  Resources->getAST().get()->runUnderLock([Pos, &Err, &Result](ParsedAST *AST) {
     if (!AST) {
       Err = llvm::make_error<llvm::StringError>("Invalid AST",
                                                 llvm::errc::invalid_argument);
       return;
     }
-    Result = clangd::findDocumentHighlights(Ctx, *AST, Pos);
+    Result = clangd::findDocumentHighlights(*AST, Pos);
   });
 
   if (Err)
     return std::move(*Err);
   return make_tagged(Result, TaggedFS.Tag);
 }
 
-std::future<Context> ClangdServer::scheduleReparseAndDiags(
-    Context Ctx, PathRef File, VersionedDraft Contents,
-    std::shared_ptr<CppFile> Resources,
+void ClangdServer::scheduleReparseAndDiags(
+    PathRef File, VersionedDraft Contents, std::shared_ptr<CppFile> Resources,
     Tagged<IntrusiveRefCntPtr<vfs::FileSystem>> TaggedFS,
     bool AllowCachedCompileFlags) {
   llvm::Optional<tooling::CompileCommand> ReusedCommand;
@@ -583,27 +553,21 @@
   ParseInputs Inputs = {std::move(Command), std::move(TaggedFS.Value),
                         *std::move(Contents.Draft)};
   assert(Contents.Draft && "Draft must have contents");
-  UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>(const Context &)>
+  UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>()>
       DeferredRebuild = Resources->deferRebuild(std::move(Inputs));
-  std::promise<Context> DonePromise;
-  std::future<Context> DoneFuture = DonePromise.get_future();
 
   DocVersion Version = Contents.Version;
   Path FileStr = File;
   VFSTag Tag = TaggedFS.Tag;
   auto ReparseAndPublishDiags =
       [this, FileStr, Version,
-       Tag](UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>(
-                const Context &)>
-                DeferredRebuild,
-            std::promise<Context> DonePromise, Context Ctx) -> void {
-    auto Guard = onScopeExit([&]() { DonePromise.set_value(std::move(Ctx)); });
-
+       Tag](UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>()>
+                DeferredRebuild) -> void {
     auto CurrentVersion = DraftMgr.getVersion(FileStr);
     if (CurrentVersion != Version)
       return; // This request is outdated
 
-    auto Diags = DeferredRebuild(Ctx);
+    auto Diags = DeferredRebuild();
     if (!Diags)
       return; // A new reparse was requested before this one completed.
 
@@ -619,37 +583,20 @@
       return;
     LastReportedDiagsVersion = Version;
 
-    DiagConsumer.onDiagnosticsReady(Ctx, FileStr,
+    DiagConsumer.onDiagnosticsReady(FileStr,
                                     make_tagged(std::move(*Diags), Tag));
   };
 
   WorkScheduler.addToFront(std::move(ReparseAndPublishDiags),
-                           std::move(DeferredRebuild), std::move(DonePromise),
-                           std::move(Ctx));
-  return DoneFuture;
+                           std::move(DeferredRebuild));
 }
 
-std::future<Context>
-ClangdServer::scheduleCancelRebuild(Context Ctx,
-                                    std::shared_ptr<CppFile> Resources) {
-  std::promise<Context> DonePromise;
-  std::future<Context> DoneFuture = DonePromise.get_future();
+void ClangdServer::scheduleCancelRebuild(std::shared_ptr<CppFile> Resources) {
   if (!Resources) {
     // No need to schedule any cleanup.
-    DonePromise.set_value(std::move(Ctx));
-    return DoneFuture;
+    return;
   }
-
-  UniqueFunction<void()> DeferredCancel = Resources->deferCancelRebuild();
-  auto CancelReparses = [Resources](std::promise<Context> DonePromise,
-                                    UniqueFunction<void()> DeferredCancel,
-                                    Context Ctx) {
-    DeferredCancel();
-    DonePromise.set_value(std::move(Ctx));
-  };
-  WorkScheduler.addToFront(std::move(CancelReparses), std::move(DonePromise),
-                           std::move(DeferredCancel), std::move(Ctx));
-  return DoneFuture;
+  WorkScheduler.addToFront(Resources->deferCancelRebuild());
 }
 
 void ClangdServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
Index: clangd/ClangdLSPServer.h
===================================================================
--- clangd/ClangdLSPServer.h
+++ clangd/ClangdLSPServer.h
@@ -50,32 +50,30 @@
 private:
   // Implement DiagnosticsConsumer.
   virtual void
-  onDiagnosticsReady(const Context &Ctx, PathRef File,
+  onDiagnosticsReady(PathRef File,
                      Tagged<std::vector<DiagWithFixIts>> Diagnostics) override;
 
   // Implement ProtocolCallbacks.
-  void onInitialize(Ctx C, InitializeParams &Params) override;
-  void onShutdown(Ctx C, ShutdownParams &Params) override;
-  void onExit(Ctx C, ExitParams &Params) override;
-  void onDocumentDidOpen(Ctx C, DidOpenTextDocumentParams &Params) override;
-  void onDocumentDidChange(Ctx C, DidChangeTextDocumentParams &Params) override;
-  void onDocumentDidClose(Ctx C, DidCloseTextDocumentParams &Params) override;
+  void onInitialize(InitializeParams &Params) override;
+  void onShutdown(ShutdownParams &Params) override;
+  void onExit(ExitParams &Params) override;
+  void onDocumentDidOpen(DidOpenTextDocumentParams &Params) override;
+  void onDocumentDidChange(DidChangeTextDocumentParams &Params) override;
+  void onDocumentDidClose(DidCloseTextDocumentParams &Params) override;
   void
-  onDocumentOnTypeFormatting(Ctx C,
-                             DocumentOnTypeFormattingParams &Params) override;
+  onDocumentOnTypeFormatting(DocumentOnTypeFormattingParams &Params) override;
   void
-  onDocumentRangeFormatting(Ctx C,
-                            DocumentRangeFormattingParams &Params) override;
-  void onDocumentFormatting(Ctx C, DocumentFormattingParams &Params) override;
-  void onCodeAction(Ctx C, CodeActionParams &Params) override;
-  void onCompletion(Ctx C, TextDocumentPositionParams &Params) override;
-  void onSignatureHelp(Ctx C, TextDocumentPositionParams &Params) override;
-  void onGoToDefinition(Ctx C, TextDocumentPositionParams &Params) override;
-  void onSwitchSourceHeader(Ctx C, TextDocumentIdentifier &Params) override;
-  void onDocumentHighlight(Ctx C, TextDocumentPositionParams &Params) override;
-  void onFileEvent(Ctx C, DidChangeWatchedFilesParams &Params) override;
-  void onCommand(Ctx C, ExecuteCommandParams &Params) override;
-  void onRename(Ctx C, RenameParams &Parames) override;
+  onDocumentRangeFormatting(DocumentRangeFormattingParams &Params) override;
+  void onDocumentFormatting(DocumentFormattingParams &Params) override;
+  void onCodeAction(CodeActionParams &Params) override;
+  void onCompletion(TextDocumentPositionParams &Params) override;
+  void onSignatureHelp(TextDocumentPositionParams &Params) override;
+  void onGoToDefinition(TextDocumentPositionParams &Params) override;
+  void onSwitchSourceHeader(TextDocumentIdentifier &Params) override;
+  void onDocumentHighlight(TextDocumentPositionParams &Params) override;
+  void onFileEvent(DidChangeWatchedFilesParams &Params) override;
+  void onCommand(ExecuteCommandParams &Params) override;
+  void onRename(RenameParams &Parames) override;
 
   std::vector<TextEdit> getFixIts(StringRef File, const clangd::Diagnostic &D);
 
Index: clangd/ClangdLSPServer.cpp
===================================================================
--- clangd/ClangdLSPServer.cpp
+++ clangd/ClangdLSPServer.cpp
@@ -45,8 +45,8 @@
 
 } // namespace
 
-void ClangdLSPServer::onInitialize(Ctx C, InitializeParams &Params) {
-  reply(C, json::obj{
+void ClangdLSPServer::onInitialize(InitializeParams &Params) {
+  reply(json::obj{
       {{"capabilities",
         json::obj{
             {"textDocumentSync", 1},
@@ -81,38 +81,35 @@
     Server.setRootPath(*Params.rootPath);
 }
 
-void ClangdLSPServer::onShutdown(Ctx C, ShutdownParams &Params) {
+void ClangdLSPServer::onShutdown(ShutdownParams &Params) {
   // Do essentially nothing, just say we're ready to exit.
   ShutdownRequestReceived = true;
-  reply(C, nullptr);
+  reply(nullptr);
 }
 
-void ClangdLSPServer::onExit(Ctx C, ExitParams &Params) { IsDone = true; }
+void ClangdLSPServer::onExit(ExitParams &Params) { IsDone = true; }
 
-void ClangdLSPServer::onDocumentDidOpen(Ctx C,
-                                        DidOpenTextDocumentParams &Params) {
+void ClangdLSPServer::onDocumentDidOpen(DidOpenTextDocumentParams &Params) {
   if (Params.metadata && !Params.metadata->extraFlags.empty())
     CDB.setExtraFlagsForFile(Params.textDocument.uri.file,
                              std::move(Params.metadata->extraFlags));
-  Server.addDocument(std::move(C), Params.textDocument.uri.file,
-                     Params.textDocument.text);
+  Server.addDocument(Params.textDocument.uri.file, Params.textDocument.text);
 }
 
-void ClangdLSPServer::onDocumentDidChange(Ctx C,
-                                          DidChangeTextDocumentParams &Params) {
+void ClangdLSPServer::onDocumentDidChange( DidChangeTextDocumentParams &Params) {
   if (Params.contentChanges.size() != 1)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "can only apply one change at a time");
   // We only support full syncing right now.
-  Server.addDocument(std::move(C), Params.textDocument.uri.file,
+  Server.addDocument(Params.textDocument.uri.file,
                      Params.contentChanges[0].text);
 }
 
-void ClangdLSPServer::onFileEvent(Ctx C, DidChangeWatchedFilesParams &Params) {
+void ClangdLSPServer::onFileEvent(DidChangeWatchedFilesParams &Params) {
   Server.onFileEvent(Params);
 }
 
-void ClangdLSPServer::onCommand(Ctx C, ExecuteCommandParams &Params) {
+void ClangdLSPServer::onCommand(ExecuteCommandParams &Params) {
   if (Params.command == ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND &&
       Params.workspaceEdit) {
     // The flow for "apply-fix" :
@@ -126,100 +123,98 @@
 
     ApplyWorkspaceEditParams ApplyEdit;
     ApplyEdit.edit = *Params.workspaceEdit;
-    reply(C, "Fix applied.");
+    reply("Fix applied.");
     // We don't need the response so id == 1 is OK.
     // Ideally, we would wait for the response and if there is no error, we
     // would reply success/failure to the original RPC.
-    call(C, "workspace/applyEdit", ApplyEdit);
+    call("workspace/applyEdit", ApplyEdit);
   } else {
     // We should not get here because ExecuteCommandParams would not have
     // parsed in the first place and this handler should not be called. But if
     // more commands are added, this will be here has a safe guard.
     replyError(
-        C, ErrorCode::InvalidParams,
+        ErrorCode::InvalidParams,
         llvm::formatv("Unsupported command \"{0}\".", Params.command).str());
   }
 }
 
-void ClangdLSPServer::onRename(Ctx C, RenameParams &Params) {
+void ClangdLSPServer::onRename(RenameParams &Params) {
   auto File = Params.textDocument.uri.file;
   auto Code = Server.getDocument(File);
   if (!Code)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "onRename called for non-added file");
 
-  auto Replacements = Server.rename(C, File, Params.position, Params.newName);
+  auto Replacements = Server.rename(File, Params.position, Params.newName);
   if (!Replacements) {
-    replyError(C, ErrorCode::InternalError,
+    replyError(ErrorCode::InternalError,
                llvm::toString(Replacements.takeError()));
     return;
   }
 
   std::vector<TextEdit> Edits = replacementsToEdits(*Code, *Replacements);
   WorkspaceEdit WE;
   WE.changes = {{Params.textDocument.uri.uri, Edits}};
-  reply(C, WE);
+  reply(WE);
 }
 
-void ClangdLSPServer::onDocumentDidClose(Ctx C,
-                                         DidCloseTextDocumentParams &Params) {
-  Server.removeDocument(std::move(C), Params.textDocument.uri.file);
+void ClangdLSPServer::onDocumentDidClose(DidCloseTextDocumentParams &Params) {
+  Server.removeDocument(Params.textDocument.uri.file);
 }
 
 void ClangdLSPServer::onDocumentOnTypeFormatting(
-    Ctx C, DocumentOnTypeFormattingParams &Params) {
+    DocumentOnTypeFormattingParams &Params) {
   auto File = Params.textDocument.uri.file;
   auto Code = Server.getDocument(File);
   if (!Code)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "onDocumentOnTypeFormatting called for non-added file");
 
   auto ReplacementsOrError = Server.formatOnType(*Code, File, Params.position);
   if (ReplacementsOrError)
-    reply(C, json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
+    reply(json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
   else
-    replyError(C, ErrorCode::UnknownErrorCode,
+    replyError(ErrorCode::UnknownErrorCode,
                llvm::toString(ReplacementsOrError.takeError()));
 }
 
 void ClangdLSPServer::onDocumentRangeFormatting(
-    Ctx C, DocumentRangeFormattingParams &Params) {
+    DocumentRangeFormattingParams &Params) {
   auto File = Params.textDocument.uri.file;
   auto Code = Server.getDocument(File);
   if (!Code)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "onDocumentRangeFormatting called for non-added file");
 
   auto ReplacementsOrError = Server.formatRange(*Code, File, Params.range);
   if (ReplacementsOrError)
-    reply(C, json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
+    reply(json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
   else
-    replyError(C, ErrorCode::UnknownErrorCode,
+    replyError(ErrorCode::UnknownErrorCode,
                llvm::toString(ReplacementsOrError.takeError()));
 }
 
-void ClangdLSPServer::onDocumentFormatting(Ctx C,
-                                           DocumentFormattingParams &Params) {
+void ClangdLSPServer::onDocumentFormatting(DocumentFormattingParams &Params) {
   auto File = Params.textDocument.uri.file;
   auto Code = Server.getDocument(File);
   if (!Code)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "onDocumentFormatting called for non-added file");
 
   auto ReplacementsOrError = Server.formatFile(*Code, File);
   if (ReplacementsOrError)
-    reply(C, json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
+    reply(json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
   else
-    replyError(C, ErrorCode::UnknownErrorCode,
+    replyError(ErrorCode::UnknownErrorCode,
                llvm::toString(ReplacementsOrError.takeError()));
 }
 
-void ClangdLSPServer::onCodeAction(Ctx C, CodeActionParams &Params) {
+void ClangdLSPServer::onCodeAction(CodeActionParams &Params) {
   // We provide a code action for each diagnostic at the requested location
   // which has FixIts available.
   auto Code = Server.getDocument(Params.textDocument.uri.file);
   if (!Code)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "onCodeAction called for non-added file");
 
   json::ary Commands;
@@ -235,68 +230,54 @@
       });
     }
   }
-  reply(C, std::move(Commands));
+  reply(std::move(Commands));
 }
 
-void ClangdLSPServer::onCompletion(Ctx C, TextDocumentPositionParams &Params) {
-  auto Reply = Server
-                   .codeComplete(std::move(C), Params.textDocument.uri.file,
-                                 Position{Params.position.line,
-                                          Params.position.character},
-                                 CCOpts)
-                   .get(); // FIXME(ibiryukov): This could be made async if we
-                           // had an API that would allow to attach callbacks to
-                           // futures returned by ClangdServer.
-
-  // We have std::move'd from C, now restore it from response of codeComplete.
-  C = std::move(Reply.first);
-  auto List = std::move(Reply.second.Value);
-  reply(C, List);
+void ClangdLSPServer::onCompletion(TextDocumentPositionParams &Params) {
+  Server.codeComplete(Params.textDocument.uri.file,
+                      Position{Params.position.line, Params.position.character},
+                      CCOpts,
+                      [](Tagged<CompletionList> List) { reply(List.Value); });
 }
 
-void ClangdLSPServer::onSignatureHelp(Ctx C,
-                                      TextDocumentPositionParams &Params) {
+void ClangdLSPServer::onSignatureHelp(TextDocumentPositionParams &Params) {
   auto SignatureHelp = Server.signatureHelp(
-      C, Params.textDocument.uri.file,
+      Params.textDocument.uri.file,
       Position{Params.position.line, Params.position.character});
   if (!SignatureHelp)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       llvm::toString(SignatureHelp.takeError()));
-  reply(C, SignatureHelp->Value);
+  reply(SignatureHelp->Value);
 }
 
-void ClangdLSPServer::onGoToDefinition(Ctx C,
-                                       TextDocumentPositionParams &Params) {
+void ClangdLSPServer::onGoToDefinition(TextDocumentPositionParams &Params) {
   auto Items = Server.findDefinitions(
-      C, Params.textDocument.uri.file,
+      Params.textDocument.uri.file,
       Position{Params.position.line, Params.position.character});
   if (!Items)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       llvm::toString(Items.takeError()));
-  reply(C, json::ary(Items->Value));
+  reply(json::ary(Items->Value));
 }
 
-void ClangdLSPServer::onSwitchSourceHeader(Ctx C,
-                                           TextDocumentIdentifier &Params) {
+void ClangdLSPServer::onSwitchSourceHeader(TextDocumentIdentifier &Params) {
   llvm::Optional<Path> Result = Server.switchSourceHeader(Params.uri.file);
   std::string ResultUri;
-  reply(C, Result ? URI::fromFile(*Result).uri : "");
+  reply(Result ? URI::fromFile(*Result).uri : "");
 }
 
-void ClangdLSPServer::onDocumentHighlight(Ctx C,
-                                          TextDocumentPositionParams &Params) {
-
+void ClangdLSPServer::onDocumentHighlight(TextDocumentPositionParams &Params) {
   auto Highlights = Server.findDocumentHighlights(
-      C, Params.textDocument.uri.file,
+      Params.textDocument.uri.file,
       Position{Params.position.line, Params.position.character});
 
   if (!Highlights) {
-    replyError(C, ErrorCode::InternalError,
+    replyError(ErrorCode::InternalError,
                llvm::toString(Highlights.takeError()));
     return;
   }
 
-  reply(C, json::ary(Highlights->Value));
+  reply(json::ary(Highlights->Value));
 }
 
 ClangdLSPServer::ClangdLSPServer(JSONOutput &Out, unsigned AsyncThreadsCount,
@@ -315,8 +296,8 @@
   assert(!IsDone && "Run was called before");
 
   // Set up JSONRPCDispatcher.
-  JSONRPCDispatcher Dispatcher([](Context Ctx, const json::Expr &Params) {
-    replyError(Ctx, ErrorCode::MethodNotFound, "method not found");
+  JSONRPCDispatcher Dispatcher([](const json::Expr &Params) {
+    replyError(ErrorCode::MethodNotFound, "method not found");
   });
   registerCallbackHandlers(Dispatcher, Out, /*Callbacks=*/*this);
 
@@ -346,8 +327,7 @@
 }
 
 void ClangdLSPServer::onDiagnosticsReady(
-    const Context &Ctx, PathRef File,
-    Tagged<std::vector<DiagWithFixIts>> Diagnostics) {
+    PathRef File, Tagged<std::vector<DiagWithFixIts>> Diagnostics) {
   json::ary DiagnosticsJSON;
 
   DiagnosticToReplacementMap LocalFixIts; // Temporary storage
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
http://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to