[clangd] Improve the tracing API

[ilya-biryukov]

Summary

The new API allows tracing events to outlive individual Spans and
gives a guarantee that invididual Spans will be not be called from
multiple threads.

[ilya-biryukov]

• Remove redundant code

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;
 } // namespace
 
 void JSONOutput::writeMessage(const json::Expr &Message) {
   std::string S;
   llvm::raw_string_ostream OS(S);
   if (Pretty)

 
 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);
+  trace::Span Span(Ctx, "clangd::reply");
+  SPAN_ATTACH(Span, "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()}}));
+
+  trace::Span Span(Ctx, "clangd::replyError");
+  SPAN_ATTACH(Span, "Error",
 
   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}}));
+
+  trace::Span Span(Ctx, "clangd::call");
+  SPAN_ATTACH(Span, "Call",
+
   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 {
   // Params should be given, use null if not.
   json::Expr Params = nullptr;
   if (auto *P = Object->get("params"))
     Params = std::move(*P);
 
   auto I = Handlers.find(*Method);
   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);
-  if (ID)
-    SPAN_ATTACH(*Tracer, "ID", *ID);
-  SPAN_ATTACH(*Tracer, "Params", Params);
+  trace::Span Span(Context::empty(), *Method);
+  SPAN_ATTACH(Span, "Params", Params);
 
-  // Update Ctx to include Tracer.
-  Ctx = std::move(Ctx).derive(RequestSpan, std::move(Tracer));
+  // Cloning traceCtx() will prolong the lifetime of Span event.
+  Context Ctx = Span.traceCtx().derive(RequestOut, &Out);
+  // Create a Context that contains request information.
+  if (ID) {
+    SPAN_ATTACH(Span, "ID", *ID);
+    Ctx = std::move(Ctx).derive(RequestID, *ID);
+  }
 
   Handler(std::move(Ctx), std::move(Params));
   return true;
 }
 
 void clangd::runLanguageServerLoop(std::istream &In, JSONOutput &Out,
                                    JSONRPCDispatcher &Dispatcher,
                                    bool &IsDone) {

clangd/Trace.h

 #include "JSONExpr.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace clang {
 namespace clangd {
 namespace trace {
 
+class TraceEvent {
+public:
+  virtual ~TraceEvent() = default;
+  /// A callback executed when a span with duration ends. Args represent data
+  /// that was attached to the event via SPAN_ATTACH.
+  virtual void endSpan(json::obj &&Args) = 0;
+  /// Called when an event, corresponding to the created span ends. The event
+  /// may stay alive longer than Spans themselves if Span::traceCtx() was
+  /// copied.
+  virtual void endEvent() = 0;
+};
+
 /// 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,
+  virtual std::unique_ptr<TraceEvent> beginSpan(const Context &Ctx,
                                                 llvm::StringRef Name) = 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

 /// 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);
+
+  // Spans are non-movable and non-copyable. If you want to prolong event
+  // defined by Span, copy the context returned by traceCtx().
+  Span(Span &&) = delete;
+  Span &operator=(Span &&) = delete;
+
+  Span(const Span &) = delete;
+  Span &operator=(const Span &) = delete;
+
   ~Span();
 
+  /// Returns a context, corresponding the the lifetime of the event started by
+  /// Span. Copies of this context will keep alive the event corresponding to
+  /// the Span.
+  Context const &traceCtx();
+
   /// 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(); }
 
 private:
   std::unique_ptr<json::obj> Args;
-  EventTracer::EndEventCallback Callback;
+  Context EventCtx;
 };
 
 #define SPAN_ATTACH(S, Name, Expr)                                             \
   do {                                                                         \
     if ((S).args() != nullptr) {                                               \
       (*((S).args()))[(Name)] = (Expr);                                        \
     }                                                                          \
   } while (0)
 
 } // namespace trace
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_TRACE_H_

clangd/Trace.cpp

 namespace clangd {
 namespace trace {
 using namespace llvm;
 
 namespace {
 // The current implementation is naive: each thread writes to Out guarded by Mu.
 // Perhaps we should replace this by something that disturbs performance less.
 class JSONTracer : public EventTracer {
+  class JSONEvent : public TraceEvent {
+  public:
+    JSONEvent(JSONTracer &Tracer) : Tracer(Tracer) {}
+
+    void endSpan(json::obj &&Args) override {
+      assert(!ArgsObj && "endSpan called twice");
+      ArgsObj = json::obj{{"args", std::move(Args)}};
+    }
+
+    void endEvent() override {
+      assert(ArgsObj && "endSpan was not called before endEvent");
+      Tracer.jsonEvent("E", std::move(*ArgsObj));
+    }
+
+  private:
+    llvm::Optional<json::obj> ArgsObj;
+    JSONTracer &Tracer;
+  };
+
 public:
   JSONTracer(raw_ostream &Out, bool Pretty)
       : Out(Out), Sep(""), Start(std::chrono::system_clock::now()),
         JSONFormat(Pretty ? "{0:2}" : "{0}") {
     // The displayTimeUnit must be ns to avoid low-precision overlap
     // calculations!
     Out << R"({"displayTimeUnit":"ns","traceEvents":[)"
         << "\n";
     rawEvent("M", json::obj{
                       {"name", "process_name"},
                       {"args", json::obj{{"name", "clangd"}}},
                   });
   }
 
   ~JSONTracer() {
     Out << "\n]}";
     Out.flush();
   }
 
-  EndEventCallback beginSpan(const Context &Ctx,
+  std::unique_ptr<TraceEvent> beginSpan(const Context &Ctx,
                                         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)}});
-    };
+    return llvm::make_unique<JSONEvent>(*this);
   }
 
   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()}});
 }
 
+namespace {
+// FIXME: we can remove the class and TraceEvent::endEvent, replace with
+// TraceEvent's destructor instead.
+class EventEnvelope {
+public:
+  EventEnvelope(std::unique_ptr<TraceEvent> Event) : Event(std::move(Event)) {
+    assert(this->Event);
+  }
+
+  EventEnvelope(EventEnvelope &&) = default;
+  EventEnvelope &operator=(EventEnvelope &&) = default;
+
+  ~EventEnvelope() {
+    if (!Event)
+      return;
+    Event->endEvent();
+  }
+
+  TraceEvent &getEvent() const { return *Event; }
+
+private:
+  std::unique_ptr<TraceEvent> Event;
+};
+
+static Key<EventEnvelope> EnvelopeKey;
+} // namespace
+
 Span::Span(const Context &Ctx, llvm::StringRef Name) {
   if (!T)
     return;
 
-  Callback = T->beginSpan(Ctx, Name);
-  if (!Callback)
+  std::unique_ptr<TraceEvent> Ev = T->beginSpan(Ctx, Name);
+  if (!Ev)
     return;
 
+  EventCtx = Ctx.derive(EnvelopeKey, EventEnvelope(std::move(Ev)));
   Args = llvm::make_unique<json::obj>();
 }
 
+Context const &Span::traceCtx() { return EventCtx; }
+
 Span::~Span() {
-  if (!Callback)
+  auto *Envelope = EventCtx.get(EnvelopeKey);
+  if (!Envelope)
     return;
 
-  assert(Args && "Args must be non-null if Callback is defined");
-  Callback(std::move(*Args));
+  assert(Args && "Args must be non-null if TraceEvent was created");
+  Envelope->getEvent().endSpan(std::move(*Args));
 }
 
 } // namespace trace
 } // namespace clangd
 } // namespace clang