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

[sammccall]

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.

[sammccall]

Remove obsolete typedef

[ilya-biryukov]

clangd/Context.h

163 ↗

(On Diff #131335)

s/const &&/&&

165 ↗

(On Diff #131335)

This will call the const& version of derive(Key, Value).
To move we need to do std::move(*this).derive(Private, std::forward<Type>(Value))

172 ↗

(On Diff #131335)

I wonder what are the use-cases for this operator ==?

clangd/Function.h

147 ↗

(On Diff #131335)

Why do we need unique_ptr instead of Optional here?

clangd/Trace.h

92 ↗

(On Diff #131335)

This is too easy to misuse unintentionally. Will this go away if we move to thread-local contexts? I.e. will the API leave out mutations of the context that can be used concurrently

[ilya-biryukov]

I've tried experimenting with that and came up with an alternative that has less changes to the API. Could you take a look, too? D42524

[sammccall]

In D42499#987583, @ilya-biryukov wrote:

I've tried experimenting with that and came up with an alternative that has less changes to the API. Could you take a look, too? D42524

As discussed offline, I've tried to incorporate the big improvements from there:

• moved back to Span objects, that for now expose the derived contexts as in your patch (and later will switch the TLS context for their lifetime)
• SPAN_ATTACH doesn't support the weird lookup that JSONRPCDispatcher uses. Rather than changing that class's traces, I've added locks to support that pattern locally.

Also as discussed offline, The proposed TraceEvent's endSpan/endEvent distinction doesn't seem necessary.
The EventTracer needs to be able to:

• store some data until the span ends
• get notified when the span ends

UniqueFunction does provide this, but the interface is slightly confusing: operator() and ~UniqueFunction both denote "when the span ends".
An object with a destructor seems a more natural way to express this, so I've kept that change - I don't feel really strongly though, happy to change this to something else simple.

clangd/Context.h

165 ↗

(On Diff #131335)

Oops, good catch!

172 ↗

(On Diff #131335)

Sorry, this was dead code from a previous iteration. Removed

clangd/Function.h

147 ↗

(On Diff #131335)

ScopeExitGuard was broken by the recent Optional changes (when Func is trivially copyable).
This is a minimal workaround, let's discuss on the other patch.

clangd/Trace.h

92 ↗

(On Diff #131335)

This is too easy to misuse unintentionally.

Agreed. I've restored the Span object, and restored the requirement that you pass the span itself to SPAN_ATTACH.

This disallows "action at a distance" that can make for nice traces e.g. for request, but can lead to thread-unsafety.
The only place where we actually use this is JSONRPCDispatcher, so I made that do it explicitly using a lock.
It seems mostly reasonable - if you feel we should remove this feature instead to avoid the complexity, happy to discuss that, but at least it's not in the main tracer interface now.

Will this go away if we move to thread-local contexts? I.e. will the API leave out mutations of the context that can be used concurrently

You've convinced me that requiring the user to pass SPAN_ATTACH to the span makes sense and we should keep it even once we have implicit context passing.

[sammccall]

Address review comments

[ilya-biryukov]

LGTM. And sorry for confusion with D42524, if any.

clangd/Function.h

147 ↗

(On Diff #131335)

LG, thanks. Wasn't aware of the breakage when writing the comment.

clangd/JSONRPCDispatcher.cpp

31 ↗

(On Diff #131460)

We could make the fields private and expose only the function to modify args under lock (similar to FillRequestInfo, i.e. the last two line of it).
That would make the class safer to use. But that's work and the class itself is an implementation detail, so we could also leave it as is.

clangd/Trace.cpp

138 ↗

(On Diff #131460)

Would be nice for Spans to have zero cost when no tracer is active. We could probably achieve that be not storing the kSpanArgs when T is null. WDYT?
That's really not a big deal, so I don't think we should block the patch on this.

[sammccall]

clangd/JSONRPCDispatcher.cpp

31 ↗

(On Diff #131460)

Yeah, good point. Encapsulated the whole mechanism in this class.

clangd/Trace.cpp

138 ↗

(On Diff #131460)

Good catch! This was left over from an older revision where SPAN_ATTACH looked up the args in the context.
Now this is for ownership only, so I skip derive() if there's no tracer, and made the key anonymous.

[sammccall]

Address review comments

clang-tools-extra/trunk/clangd/ClangdUnit.cpp

auto DoRebuildPreamble =
           /*StoreInMemory=*/That->StorePreamblesInMemory,
           SerializedDeclsCollector);
 
       // When building the AST for the main file, we do want the function
       // bodies.
       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;
       }
     };
 
     // Compute updated Preamble.
     std::shared_ptr<const PreambleData> NewPreamble = DoRebuildPreamble();
     // Publish the new Preamble.
if (NewPreamble) {
                          NewPreamble->Diags.end());
     }
 
     // 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),
-                                std::move(ContentsBuffer), PCHs, Inputs.FS);
+      NewAST =
+          ParsedAST::Build(Tracer.Ctx, std::move(CI), std::move(NewPreamble),
     }
 
     if (NewAST) {
       Diagnostics.insert(Diagnostics.end(), NewAST->getDiagnostics().begin(),
                          NewAST->getDiagnostics().end());
       if (That->ASTCallback)
         That->ASTCallback(Ctx, That->FileName, NewAST.getPointer());
     } else {

clang-tools-extra/trunk/clangd/Context.h

 ///    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.
+///
+/// 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
   /// functions that require a context when no explicit context is available.
   static Context empty();
 
 private:
   struct Data;
public:
   derive(const Key<Type> &Key,
          typename std::decay<Type>::type Value) && /* takes ownership */ {
     return Context(std::make_shared<Data>(Data{
         /*Parent=*/std::move(DataPtr), &Key,
         llvm::make_unique<TypedAnyStorage<typename std::decay<Type>::type>>(
             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;
 
 private:
   class AnyStorage {
   public:
     virtual ~AnyStorage() = default;
     virtual void *getValuePtr() = 0;

clang-tools-extra/trunk/clangd/JSONRPCDispatcher.cpp

 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/SourceMgr.h"
 #include <istream>
 
 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) {
   std::string S;
   llvm::raw_string_ostream OS(S);
   if (Pretty)
     OS << llvm::formatv("{0:2}", Message);
   else

 
 void clangd::reply(const Context &Ctx, json::Expr &&Result) {
   auto ID = Ctx.get(RequestID);
   if (!ID) {
     log(Ctx, "Attempted to reply to a notification!");
     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{
           {"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()}}));
+  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)
         ->writeMessage(json::obj{
             {"jsonrpc", "2.0"},
             {"id", *ID},
             {"error",
              json::obj{{"code", static_cast<int>(code)}, {"message", Message}}},
         });
   }
 }
 
 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"},
           {"id", 1},
           {"method", Method},
           {"params", std::move(Params)},
       });
 }
bool JSONRPCDispatcher::call(const json::Expr &Message, JSONOutput &Out) const {
   auto &Handler = I != Handlers.end() ? I->second : UnknownHandler;
 
   // Create a Context that contains request information.
   auto Ctx = Context::empty().derive(RequestOut, &Out);
   if (ID)
     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;
 }
 
 void clangd::runLanguageServerLoop(std::istream &In, JSONOutput &Out,
                                    JSONRPCDispatcher &Dispatcher,
                                    bool &IsDone) {
   while (In.good()) {
     // A Language Server Protocol message starts with a set of HTTP headers,

clang-tools-extra/trunk/clangd/Trace.h

 namespace clang {
 namespace clangd {
 namespace trace {
 
 /// A consumer of trace events. The events are produced by Spans and trace::log.
 /// 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,
                        json::obj &&Args) = 0;
 };
 
 /// Sets up a global EventTracer that consumes events produced by Span and
 /// trace::log. Only one TracingSession can be active at a time and it should be

 /// https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview
 std::unique_ptr<EventTracer> createJSONTracer(llvm::raw_ostream &OS,
                                               bool Pretty = false);
 
 /// Records a single instant event, associated with the current thread.
 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
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_TRACE_H_

clang-tools-extra/trunk/clangd/Trace.cpp

 //===--- Trace.cpp - Performance tracing facilities -----------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "llvm/Support/Threading.h"
 #include <mutex>
 
 namespace clang {
 namespace clangd {
rawEvent("M", json::obj{
                   });
   }
 
   ~JSONTracer() {
     Out << "\n]}";
     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,
                json::obj &&Args) override {
     jsonEvent("i", json::obj{{"name", Name}, {"args", std::move(Args)}});
   }
 
   // Record an event on the current thread. ph, pid, tid, ts are set.

 }
 
 void log(const Context &Ctx, const Twine &Message) {
   if (!T)
     return;
   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
 } // namespace clang

clang-tools-extra/trunk/unittests/clangd/TraceTests.cpp

 TEST(TraceTest, SmokeTest) {
   // Capture some events.
   std::string JSON;
   {
     raw_string_ostream OS(JSON);
     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");
     }
   }
 
   // Get the root JSON object using the YAML parser.
   SourceMgr SM;
   yaml::Stream Stream(JSON, SM);
   auto Doc = Stream.begin();
   ASSERT_NE(Doc, Stream.end());