[clangd] Modify the Span API so that Spans propagate with contexts.

Summary:
This is probably the right behavior for distributed tracers, and makes unpaired
begin-end events impossible without requiring Spans to be bound to a thread.

The API is conceptually clean but syntactically awkward. As discussed offline,
this is basically a naming problem and will go away if (when) we use TLS to
store the current context.

The apparently-unrelated change to onScopeExit are because its move semantics
broken if Func is POD-like since r322838. This is true of function pointers,
and the lambda I use here that captures two pointers only.
I've raised this issue on llvm-dev and will revert this part if we fix it in
some other way.

Reviewers: ilya-biryukov

Subscribers: klimek, jkorous-apple, ioeric, cfe-commits

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

6 files changed, 102 insertions, 68 deletions

diff --git a/clang-tools-extra/clangd/ClangdUnit.cpp b/clang-tools-extra/clangd/ClangdUnit.cpp
index ed06607ef3c..d915871a4ef 100644
--- a/clang-tools-extra/clangd/ClangdUnit.cpp
+++ b/clang-tools-extra/clangd/ClangdUnit.cpp
@@ -564,15 +564,16 @@ CppFile::deferRebuild(ParseInputs &&Inputs) {
       CI->getFrontendOpts().SkipFunctionBodies = false;
 
       if (BuiltPreamble) {
-        log(Ctx, "Built preamble of size " + Twine(BuiltPreamble->getSize()) +
-                     " for file " + Twine(That->FileName));
+        log(Tracer.Ctx, "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,
+        log(Tracer.Ctx,
             "Could not build a preamble for file " + Twine(That->FileName));
         return nullptr;
       }
@@ -605,8 +606,9 @@ CppFile::deferRebuild(ParseInputs &&Inputs) {
     {
       trace::Span Tracer(Ctx, "Build");
       SPAN_ATTACH(Tracer, "File", That->FileName);
-      NewAST = ParsedAST::Build(Ctx, std::move(CI), std::move(NewPreamble),
-                                std::move(ContentsBuffer), PCHs, Inputs.FS);
+      NewAST =
+          ParsedAST::Build(Tracer.Ctx, std::move(CI), std::move(NewPreamble),
+                           std::move(ContentsBuffer), PCHs, Inputs.FS);
     }
 
     if (NewAST) {
diff --git a/clang-tools-extra/clangd/Context.h b/clang-tools-extra/clangd/Context.h
index c64780caf44..5254b04e110 100644
--- a/clang-tools-extra/clangd/Context.h
+++ b/clang-tools-extra/clangd/Context.h
@@ -84,6 +84,9 @@ public:
 ///
 /// Keys are typically used across multiple functions, so most of the time you
 /// would want to make them static class members or global variables.
+///
+/// FIXME: Rather than manual plumbing, pass Context using thread-local storage
+/// by default, and make thread boundaries deal with propagation explicitly.
 class Context {
 public:
   /// Returns an empty context that contains no data. Useful for calling
@@ -150,6 +153,18 @@ public:
             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) && {
+    static Key<typename std::decay<Type>::type> Private;
+    return std::move(this)->derive(Private, std::forward<Type>(Value));
+  }
+
   /// Clone this context object.
   Context clone() const;
 
diff --git a/clang-tools-extra/clangd/JSONRPCDispatcher.cpp b/clang-tools-extra/clangd/JSONRPCDispatcher.cpp
index fec7cd7ec7c..f0598493932 100644
--- a/clang-tools-extra/clangd/JSONRPCDispatcher.cpp
+++ b/clang-tools-extra/clangd/JSONRPCDispatcher.cpp
@@ -20,9 +20,35 @@ using namespace clang;
 using namespace clangd;
 
 namespace {
-static Key<std::unique_ptr<trace::Span>> RequestSpan;
 static Key<json::Expr> RequestID;
 static Key<JSONOutput *> RequestOut;
+
+// When tracing, we trace a request and attach the repsonse in reply().
+// Because the Span isn't available, we find the current request using Context.
+class RequestSpan {
+  RequestSpan(json::obj *Args) : Args(Args) {}
+  std::mutex Mu;
+  json::obj *Args;
+  static Key<std::unique_ptr<RequestSpan>> Key;
+
+public:
+  // Return a context that's aware of the enclosing request, identified by Span.
+  static Context stash(const trace::Span &Span) {
+    return Span.Ctx.derive(RequestSpan::Key, std::unique_ptr<RequestSpan>(
+                                                 new RequestSpan(Span.Args)));
+  }
+
+  // If there's an enclosing request and the tracer is interested, calls \p F
+  // with a json::obj where request info can be added.
+  template <typename Func> static void attach(const Context &Ctx, Func &&F) {
+    auto *RequestArgs = Ctx.get(RequestSpan::Key);
+    if (!RequestArgs || !*RequestArgs || !(*RequestArgs)->Args)
+      return;
+    std::lock_guard<std::mutex> Lock((*RequestArgs)->Mu);
+    F(*(*RequestArgs)->Args);
+  }
+};
+Key<std::unique_ptr<RequestSpan>> RequestSpan::Key;
 } // namespace
 
 void JSONOutput::writeMessage(const json::Expr &Message) {
@@ -66,8 +92,7 @@ void clangd::reply(const Context &Ctx, json::Expr &&Result) {
     return;
   }
 
-  if (auto *Span = Ctx.get(RequestSpan))
-    SPAN_ATTACH(**Span, "Reply", Result);
+  RequestSpan::attach(Ctx, [&](json::obj &Args) { Args["Reply"] = Result; });
 
   Ctx.getExisting(RequestOut)
       ->writeMessage(json::obj{
@@ -80,10 +105,10 @@ void clangd::reply(const Context &Ctx, json::Expr &&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()}}));
+  RequestSpan::attach(Ctx, [&](json::obj &Args) {
+    Args["Error"] =
+        json::obj{{"code", static_cast<int>(code)}, {"message", Message.str()}};
+  });
 
   if (auto ID = Ctx.get(RequestID)) {
     Ctx.getExisting(RequestOut)
@@ -99,9 +124,9 @@ void clangd::replyError(const Context &Ctx, ErrorCode code,
 void clangd::call(const Context &Ctx, 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",
-                (json::obj{{"method", Method.str()}, {"params", Params}}));
+  RequestSpan::attach(Ctx, [&](json::obj &Args) {
+    Args["Call"] = json::obj{{"method", Method.str()}, {"params", Params}};
+  });
   Ctx.getExisting(RequestOut)
       ->writeMessage(json::obj{
           {"jsonrpc", "2.0"},
@@ -143,15 +168,13 @@ bool JSONRPCDispatcher::call(const json::Expr &Message, JSONOutput &Out) const {
     Ctx = std::move(Ctx).derive(RequestID, *ID);
 
   // Create a tracing Span covering the whole request lifetime.
-  auto Tracer = llvm::make_unique<trace::Span>(Ctx, *Method);
+  trace::Span Tracer(Ctx, *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(Tracer, "ID", *ID);
+  SPAN_ATTACH(Tracer, "Params", Params);
 
-  Handler(std::move(Ctx), std::move(Params));
+  // Stash a reference to the span args, so later calls can add metadata.
+  Handler(RequestSpan::stash(Tracer), std::move(Params));
   return true;
 }
 
diff --git a/clang-tools-extra/clangd/Trace.cpp b/clang-tools-extra/clangd/Trace.cpp
index 03b70b09c77..7c5df89f19a 100644
--- a/clang-tools-extra/clangd/Trace.cpp
+++ b/clang-tools-extra/clangd/Trace.cpp
@@ -8,7 +8,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "Trace.h"
+#include "Context.h"
+#include "Function.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Chrono.h"
 #include "llvm/Support/FormatProviders.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -43,14 +46,12 @@ public:
     Out.flush();
   }
 
-  EndEventCallback beginSpan(const Context &Ctx,
-                             llvm::StringRef Name) override {
+  Context beginSpan(const Context &Ctx, llvm::StringRef Name,
+                    json::obj *Args) 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)}});
-    };
+    return Ctx.derive(make_scope_exit([this, Args] {
+      jsonEvent("E", json::obj{{"args", std::move(*Args)}});
+    }));
   }
 
   void instant(const Context &Ctx, llvm::StringRef Name,
@@ -125,24 +126,14 @@ void log(const Context &Ctx, const Twine &Message) {
   T->instant(Ctx, "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;
-
-  Args = llvm::make_unique<json::obj>();
-}
-
-Span::~Span() {
-  if (!Callback)
-    return;
-
-  assert(Args && "Args must be non-null if Callback is defined");
-  Callback(std::move(*Args));
-}
+// Span keeps a non-owning pointer to the args, which is how users access them.
+// The args are owned by the context though. They stick around until the
+// beginSpan() context is destroyed, when the tracing engine will consume them.
+Span::Span(const Context &Ctx, llvm::StringRef Name)
+    : Args(T ? new json::obj() : nullptr),
+      Ctx(T ? T->beginSpan(Ctx.derive(std::unique_ptr<json::obj>(Args)), Name,
+                           Args)
+            : Ctx.clone()) {}
 
 } // namespace trace
 } // namespace clangd
diff --git a/clang-tools-extra/clangd/Trace.h b/clang-tools-extra/clangd/Trace.h
index 3ad3498eaa0..ad3182bd478 100644
--- a/clang-tools-extra/clangd/Trace.h
+++ b/clang-tools-extra/clangd/Trace.h
@@ -32,17 +32,15 @@ namespace trace {
 /// 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 *Args, only complete when the event ends.
+  virtual Context beginSpan(const Context &Ctx, llvm::StringRef Name,
+                            json::obj *Args) = 0;
 
   /// Called for instant events.
   virtual void instant(const Context &Ctx, llvm::StringRef Name,
@@ -70,30 +68,35 @@ std::unique_ptr<EventTracer> createJSONTracer(llvm::raw_ostream &OS,
 void log(const Context &Ctx, const llvm::Twine &Name);
 
 /// Records an event whose duration is the lifetime of the Span object.
+/// This lifetime is extended when the span's context is reused.
+///
 /// 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);
+///
 /// 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.
+  /// Mutable metadata, if this span is interested.
   /// Prefer to use SPAN_ATTACH rather than accessing this directly.
-  json::obj *args() { return Args.get(); }
-
-private:
-  std::unique_ptr<json::obj> Args;
-  EventTracer::EndEventCallback Callback;
+  json::obj *const Args;
+  /// Propagating this context will keep the span alive.
+  const Context Ctx;
 };
 
+/// Returns mutable span metadata if this span is interested.
+/// Prefer to use SPAN_ATTACH rather than accessing this directly.
+json::obj *spanArgs(const Context &Ctx);
+
+/// Attach a key-value pair to a Span event.
+/// This is not threadsafe when used with the same Span.
 #define SPAN_ATTACH(S, Name, Expr)                                             \
   do {                                                                         \
-    if ((S).args() != nullptr) {                                               \
-      (*((S).args()))[(Name)] = (Expr);                                        \
-    }                                                                          \
+    if (auto *Args = (S).Args)                                                 \
+      (*Args)[Name] = Expr;                                                    \
   } while (0)
 
 } // namespace trace
diff --git a/clang-tools-extra/unittests/clangd/TraceTests.cpp b/clang-tools-extra/unittests/clangd/TraceTests.cpp
index 0f734e9784e..06f967920d2 100644
--- a/clang-tools-extra/unittests/clangd/TraceTests.cpp
+++ b/clang-tools-extra/unittests/clangd/TraceTests.cpp
@@ -78,8 +78,8 @@ TEST(TraceTest, SmokeTest) {
     auto JSONTracer = trace::createJSONTracer(OS);
     trace::Session Session(*JSONTracer);
     {
-      trace::Span S(Context::empty(), "A");
-      trace::log(Context::empty(), "B");
+      trace::Span Tracer(Context::empty(), "A");
+      trace::log(Tracer.Ctx, "B");
     }
   }