[clangd] Implemented tracing using Context

[ilya-biryukov]

clangd/JSONRPCDispatcher.h

80

This is still wrong, Context is used by Tracer internally and should not be moved.

[ilya-biryukov]

• Updated tracing to use the new Context implementation.
• Restored global tracer.

[ilya-biryukov]

• Use getExistingKey instead of getPtr in JSONRPCDispatcher.

[ilya-biryukov]

Updated the patch after changes to Context

[ilya-biryukov]

• Use derive(key, value) instead of derive().add(key, value).

[ilya-biryukov]

Merged with head

[sammccall]

clangd/Trace.cpp

117

This is a little unfortunate. Oh, well.

clangd/Trace.h

37

hmm, trace::TracingSession seems redundant? Up to you

[ilya-biryukov]

• Renamed TracingSession to Session (it is always used qualified via trace::Session anyway)

clangd/ClangdUnit.cpp

auto DoRebuildPreamble =
       if (OldPreamble && OldPreamble->Preamble.CanReuse(
                              *CI, ContentsBuffer.get(), Bounds, VFS.get())) {
         return OldPreamble;
       }
       // 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("Preamble");
+      trace::Span Tracer(Ctx, "Preamble");
       SPAN_ATTACH(Tracer, "File", That->FileName);
       std::vector<DiagWithFixIts> PreambleDiags;
       StoreDiagsConsumer PreambleDiagnosticsConsumer(/*ref*/ PreambleDiags);
       IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
           CompilerInstance::createDiagnostics(
               &CI->getDiagnosticOpts(), &PreambleDiagnosticsConsumer, false);
       CppFilePreambleCallbacks SerializedDeclsCollector;
       auto BuiltPreamble = PrecompiledPreamble::Build(
auto FinishRebuild = [OldPreamble, VFS, RequestRebuildCounter, PCHs,
     if (NewPreamble) {
       Diagnostics.insert(Diagnostics.begin(), NewPreamble->Diags.begin(),
                          NewPreamble->Diags.end());
     }
 
     // Compute updated AST.
     llvm::Optional<ParsedAST> NewAST;
     {
-      trace::Span Tracer("Build");
+      trace::Span Tracer(Ctx, "Build");
       SPAN_ATTACH(Tracer, "File", That->FileName);
       NewAST = ParsedAST::Build(std::move(CI), std::move(NewPreamble),
                                 std::move(ContentsBuffer), PCHs, VFS, Ctx);
     }
 
     if (NewAST) {
       Diagnostics.insert(Diagnostics.end(), NewAST->getDiagnostics().begin(),
                          NewAST->getDiagnostics().end());

clangd/JSONRPCDispatcher.h

 };
 
 /// Context object passed to handlers to allow replies.
 class RequestContext {
 public:
   RequestContext(Context Ctx, JSONOutput &Out, StringRef Method,
                  llvm::Optional<json::Expr> ID)
       : Ctx(std::move(Ctx)), Out(Out), ID(std::move(ID)),
-        Tracer(llvm::make_unique<trace::Span>(Method)) {
+        Tracer(llvm::make_unique<trace::Span>(this->Ctx, Method)) {
     if (this->ID)
       SPAN_ATTACH(tracer(), "ID", *this->ID);
   }
 
   /// Sends a successful reply.
   void reply(json::Expr &&Result);
   /// Sends an error response to the client, and logs it.
   void replyError(ErrorCode code, const llvm::StringRef &Message);
   /// Sends a request to the client.
   void call(llvm::StringRef Method, json::Expr &&Params);
 
   trace::Span &tracer() { return *Tracer; }
   Context &ctx() { return Ctx; }
 
 private:
+  // Ctx must be before Tracer!
   Context Ctx;
   JSONOutput &Out;
   llvm::Optional<json::Expr> ID;
   std::unique_ptr<trace::Span> Tracer;
 };
 
 /// Main JSONRPC entry point. This parses the JSONRPC "header" and calls the
 /// registered Handler for the method received.
 class JSONRPCDispatcher {

clangd/JSONRPCDispatcher.cpp

void JSONOutput::writeMessage(const json::Expr &Message) {
   Logs.flush();
 
   // Emit message with header.
   Outs << "Content-Length: " << S.size() << "\r\n\r\n" << S;
   Outs.flush();
 }
 
 void JSONOutput::logImpl(Context &Ctx, const Twine &Message) {
-  // FIXME(ibiryukov): get rid of trace::log here.
-  trace::log(Message);
+  trace::log(Ctx, Message);
   std::lock_guard<std::mutex> Guard(StreamMutex);
   Logs << Message;
   Logs.flush();
 }
 
 void JSONOutput::mirrorInput(const Twine &Message) {
   if (!InputMirror)
     return;

clangd/ProtocolHandlers.cpp

struct HandlerRegisterer {
   template <typename Param>
   void operator()(StringRef Method,
                   void (ProtocolCallbacks::*Handler)(RequestContext, Param)) {
     // Capture pointers by value, as the lambda will outlive this object.
     auto *Callbacks = this->Callbacks;
     Dispatcher.registerHandler(
         Method, [=](RequestContext C, llvm::yaml::MappingNode *RawParams) {
           if (auto P = [&] {
-                trace::Span Tracer("Parse");
+                trace::Span Tracer(C.ctx(), "Parse");
                 return std::decay<Param>::type::parse(RawParams, C.ctx());
               }()) {
             (Callbacks->*Handler)(std::move(C), *P);
           } else {
             log(C.ctx(), "Failed to decode " + Method + " request.\n");
           }
         });
   }

clangd/Trace.h

 //===--- Trace.h - Performance tracing facilities ---------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // Supports writing performance traces describing clangd's behavior.
-// Traces are written in the Trace Event format supported by chrome's trace
-// viewer (chrome://tracing).
+// Traces are consumed by implementations of the EventTracer interface.
 //
-// The format is documented here:
-// https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview
 //
 // All APIs are no-ops unless a Session is active (created by ClangdMain).
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_TRACE_H_
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_TRACE_H_
 
+#include "Context.h"
 #include "JSONExpr.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace clang {
 namespace clangd {
 namespace trace {
 
-// A session directs the output of trace events. Only one Session can exist.
-// It should be created before clangd threads are spawned, and destroyed after
-// they exit.
-// TODO: we may want to add pluggable support for other tracing backends.
-class Session {
+class EventTracer {
 public:
-  // Starts a sessions capturing trace events and writing Trace Event JSON.
-  static std::unique_ptr<Session> create(llvm::raw_ostream &OS,
-                                         bool Pretty = false);
-  ~Session();
+  static PtrKey<EventTracer> CtxKey;
 
-private:
-  Session() = default;
+  virtual ~EventTracer() = default;
+  virtual void event(llvm::StringRef Phase, json::obj &&Contents) = 0;
 };
 
-// Records a single instant event, associated with the current thread.
-void log(const llvm::Twine &Name);
+/// Create an instance of EventTracer that produces an output in the Trace Event
+/// format supported by Chrome's trace viewer (chrome://tracing).
+///
+/// The format is documented here:
+/// https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview
+///
+/// The implementation supports concurrent calls and can be used as a global
+/// tracer (i.e., can be put into a global Context).
+std::unique_ptr<EventTracer> createJSONTracer(llvm::raw_ostream &OS,
+                                              bool Pretty = false);
 
-// Records an event whose duration is the lifetime of the Span object.
-//
-// 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.
+/// Records a single instant event, associated with the current thread.
+void log(Context &Ctx, const llvm::Twine &Name);
+
+/// Records an event whose duration is the lifetime of the Span object.
+/// 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(std::string Name);
+  Span(Context &Ctx, std::string Name);
   ~Span();
 
-  // Returns mutable span metadata if this span is interested.
-  // Prefer to use SPAN_ATTACH rather than accessing this directly.
+  /// 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:
+  Context &Ctx;
   std::unique_ptr<json::obj> Args;
 };
 
 #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 clang {
 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 Tracer {
+class JSONTracer : public EventTracer {
 public:
-  Tracer(raw_ostream &Out, bool Pretty)
+  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"}}},
                   });
   }
 
-  ~Tracer() {
+  ~JSONTracer() {
     Out << "\n]}";
     Out.flush();
   }
 
   // Record an event on the current thread. ph, pid, tid, ts are set.
   // Contents must be a list of the other JSON key/values.
-  void event(StringRef Phase, json::obj &&Contents) {
+  void event(StringRef Phase, json::obj &&Contents) override {
     uint64_t TID = get_threadid();
     std::lock_guard<std::mutex> Lock(Mu);
     // If we haven't already, emit metadata describing this thread.
     if (ThreadsWithMD.insert(TID).second) {
       SmallString<32> Name;
       get_thread_name(Name);
       if (!Name.empty()) {
         rawEvent("M", json::obj{
private:
 
   std::mutex Mu;
   raw_ostream &Out /*GUARDED_BY(Mu)*/;
   const char *Sep /*GUARDED_BY(Mu)*/;
   DenseSet<uint64_t> ThreadsWithMD /*GUARDED_BY(Mu)*/;
   const sys::TimePoint<> Start;
   const char *JSONFormat;
 };
-
-static Tracer *T = nullptr;
 } // namespace
 
-std::unique_ptr<Session> Session::create(raw_ostream &OS, bool Pretty) {
-  assert(!T && "A session is already active");
-  T = new Tracer(OS, Pretty);
-  return std::unique_ptr<Session>(new Session());
-}
+PtrKey<EventTracer> EventTracer::CtxKey;
 
-Session::~Session() {
-  delete T;
-  T = nullptr;
+std::unique_ptr<EventTracer> createJSONTracer(llvm::raw_ostream &OS,
+                                              bool Pretty) {
+  return llvm::make_unique<JSONTracer>(OS, Pretty);
 }
 
-void log(const Twine &Message) {
+void log(Context &Ctx, const Twine &Message) {
+  EventTracer *T = Ctx.get(EventTracer::CtxKey);
   if (!T)
     return;
   T->event("i", json::obj{
                     {"name", "Log"},
                     {"args", json::obj{{"Message", Message.str()}}},
                 });
 }
 
-Span::Span(std::string Name) {
+Span::Span(Context &Ctx, std::string Name) : Ctx(Ctx) {
+  EventTracer* T = Ctx.get(EventTracer::CtxKey);
   if (!T)
     return;
   T->event("B", json::obj{{"name", std::move(Name)}});
   Args = llvm::make_unique<json::obj>();
 }
 
 Span::~Span() {
+  auto T = Ctx.get(EventTracer::CtxKey);
   if (!T)
     return;
   if (!Args)
     Args = llvm::make_unique<json::obj>();
   T->event("E", Args ? json::obj{{"args", std::move(*Args)}} : json::obj{});
 }
 
 } // namespace trace
 } // namespace clangd
 } // namespace clang

clangd/tool/ClangdMain.cpp

if (!InputMirrorFile.empty()) {
     InputMirrorStream.emplace(InputMirrorFile, /*ref*/ EC, llvm::sys::fs::F_RW);
     if (EC) {
       InputMirrorStream.reset();
       llvm::errs() << "Error while opening an input mirror file: "
                    << EC.message();
     }
   }
   llvm::Optional<llvm::raw_fd_ostream> TraceStream;
-  std::unique_ptr<trace::Session> TraceSession;
+  std::unique_ptr<trace::EventTracer> TraceSession;
   if (!TraceFile.empty()) {
     std::error_code EC;
     TraceStream.emplace(TraceFile, /*ref*/ EC, llvm::sys::fs::F_RW);
     if (EC) {
       TraceFile.reset();
       llvm::errs() << "Error while opening trace file: " << EC.message();
     } else {
-      TraceSession = trace::Session::create(*TraceStream, PrettyPrint);
+      TraceSession = trace::createJSONTracer(*TraceStream, PrettyPrint);
     }
   }
 
   llvm::raw_ostream &Outs = llvm::outs();
   llvm::raw_ostream &Logs = llvm::errs();
   JSONOutput Out(Outs, Logs,
                  InputMirrorStream ? InputMirrorStream.getPointer() : nullptr,
                  PrettyPrint);
 
-  GlobalSession Session(buildCtx().add(Logger::CtxKey, &Out));
+  TypedValueMap ContextMap;
+  ContextMap.emplace(Logger::CtxKey, &Out);
+  if (TraceSession)
+    ContextMap.emplace(trace::EventTracer::CtxKey, TraceSession.get());
+  GlobalSession Session(Context(/*Parent=*/nullptr, std::move(ContextMap)));
 
   // If --compile-commands-dir arg was invoked, check value and override default
   // path.
   llvm::Optional<Path> CompileCommandsDirPath;
 
   if (CompileCommandsDir.empty()) {
     CompileCommandsDirPath = llvm::None;
   } else if (!llvm::sys::path::is_absolute(CompileCommandsDir) ||

unittests/clangd/TraceTests.cpp

 //===-- TraceTests.cpp - Tracing unit tests ---------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
+#include "Context.h"
 #include "Trace.h"
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/Threading.h"
 #include "llvm/Support/YAMLParser.h"
 #include "gmock/gmock.h"
bool VerifyObject(yaml::Node &N, std::map<std::string, std::string> Expected) {
   return Match;
 }
 
 TEST(TraceTest, SmokeTest) {
   // Capture some events.
   std::string JSON;
   {
     raw_string_ostream OS(JSON);
-    auto Session = trace::Session::create(OS);
+    auto JSONTracer = trace::createJSONTracer(OS);
+    Context TraceCtx =
+        buildCtx().add(trace::EventTracer::CtxKey, JSONTracer.get());
     {
-      trace::Span S("A");
-      trace::log("B");
+      trace::Span S(TraceCtx, "A");
+      trace::log(TraceCtx, "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());