[clangd] Implemented tracing using Context

Reviewers: sammccall, ioeric, hokein

Reviewed By: sammccall

Subscribers: klimek, luckygeck, cfe-commits

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

git-svn-id: https://llvm.org/svn/llvm-project/clang-tools-extra/trunk@320706 91177308-0d34-0410-b5e6-96231b3b80d8

6 files changed, 90 insertions, 58 deletions

diff --git a/clangd/ClangdUnit.cpp b/clangd/ClangdUnit.cpp
index e72dc3b5..5f181a0b 100644
--- a/clangd/ClangdUnit.cpp
+++ b/clangd/ClangdUnit.cpp
@@ -770,7 +770,7 @@ CppFile::deferRebuild(StringRef NewContents,
       // (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);
@@ -816,7 +816,7 @@ CppFile::deferRebuild(StringRef NewContents,
     // 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(Ctx, std::move(CI), std::move(NewPreamble),
                                 std::move(ContentsBuffer), PCHs, VFS);
diff --git a/clangd/JSONRPCDispatcher.cpp b/clangd/JSONRPCDispatcher.cpp
index e9a203d7..c67221a2 100644
--- a/clangd/JSONRPCDispatcher.cpp
+++ b/clangd/JSONRPCDispatcher.cpp
@@ -45,7 +45,7 @@ void JSONOutput::writeMessage(const json::Expr &Message) {
 }
 
 void JSONOutput::log(const Context &Ctx, const Twine &Message) {
-  trace::log(Message);
+  trace::log(Ctx, Message);
   std::lock_guard<std::mutex> Guard(StreamMutex);
   Logs << Message << '\n';
   Logs.flush();
@@ -137,16 +137,19 @@ bool JSONRPCDispatcher::call(const json::Expr &Message, JSONOutput &Out) const {
   auto I = Handlers.find(*Method);
   auto &Handler = I != Handlers.end() ? I->second : UnknownHandler;
 
-  auto Tracer = llvm::make_unique<trace::Span>(*Method);
+  // 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);
 
-  auto Ctx = Context::empty()
-                 .derive(RequestOut, &Out)
-                 .derive(RequestSpan, std::move(Tracer));
-  if (ID)
-    Ctx = std::move(Ctx).derive(RequestID, *ID);
+  // Update Ctx to include Tracer.
+  Ctx = std::move(Ctx).derive(RequestSpan, std::move(Tracer));
 
   Handler(std::move(Ctx), std::move(Params));
   return true;
diff --git a/clangd/Trace.cpp b/clangd/Trace.cpp
index 60340149..e0f09218 100644
--- a/clangd/Trace.cpp
+++ b/clangd/Trace.cpp
@@ -24,9 +24,9 @@ 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
@@ -39,14 +39,15 @@ public:
                   });
   }
 
-  ~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(const Context &Ctx, 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.
@@ -90,33 +91,38 @@ private:
   const char *JSONFormat;
 };
 
-static Tracer *T = nullptr;
+EventTracer *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());
+Session::Session(EventTracer &Tracer) {
+  assert(!T && "Resetting global tracer is not allowed.");
+  T = &Tracer;
 }
 
-Session::~Session() {
-  delete T;
-  T = nullptr;
+Session::~Session() { 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(const Context &Ctx, const Twine &Message) {
   if (!T)
     return;
-  T->event("i", json::obj{
-                    {"name", "Log"},
-                    {"args", json::obj{{"Message", Message.str()}}},
-                });
+  T->event(Ctx, "i",
+           json::obj{
+               {"name", "Log"},
+               {"args", json::obj{{"Message", Message.str()}}},
+           });
 }
 
-Span::Span(std::string Name) {
+Span::Span(const Context &Ctx, std::string Name) {
   if (!T)
     return;
-  T->event("B", json::obj{{"name", std::move(Name)}});
+  // Clone the context, so that the original Context can be moved.
+  this->Ctx.emplace(Ctx.clone());
+
+  T->event(*this->Ctx, "B", json::obj{{"name", std::move(Name)}});
   Args = llvm::make_unique<json::obj>();
 }
 
@@ -125,7 +131,8 @@ Span::~Span() {
     return;
   if (!Args)
     Args = llvm::make_unique<json::obj>();
-  T->event("E", Args ? json::obj{{"args", std::move(*Args)}} : json::obj{});
+  T->event(*Ctx, "E",
+           Args ? json::obj{{"args", std::move(*Args)}} : json::obj{});
 }
 
 } // namespace trace
diff --git a/clangd/Trace.h b/clangd/Trace.h
index 436fc836..51421bd1 100644
--- a/clangd/Trace.h
+++ b/clangd/Trace.h
@@ -8,11 +8,8 @@
 //===----------------------------------------------------------------------===//
 //
 // 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).
 //
@@ -21,6 +18,7 @@
 #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"
@@ -29,39 +27,55 @@ 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.
+/// A consumer of trace events. The events are produced by Spans and trace::log.
+class EventTracer {
+public:
+  virtual ~EventTracer() = default;
+  /// Consume a trace event.
+  virtual void event(const Context &Ctx, llvm::StringRef Phase,
+                     json::obj &&Contents) = 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
+/// set up before calling any clangd-specific functions.
 class Session {
 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(EventTracer &Tracer);
   ~Session();
-
-private:
-  Session() = default;
 };
 
-// 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(const 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(const 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:
+  llvm::Optional<Context> Ctx;
   std::unique_ptr<json::obj> Args;
 };
 
diff --git a/clangd/tool/ClangdMain.cpp b/clangd/tool/ClangdMain.cpp
index d1e4ed93..4c91fe86 100644
--- a/clangd/tool/ClangdMain.cpp
+++ b/clangd/tool/ClangdMain.cpp
@@ -115,8 +115,10 @@ int main(int argc, char *argv[]) {
                    << EC.message();
     }
   }
+
+  // Setup tracing facilities.
   llvm::Optional<llvm::raw_fd_ostream> TraceStream;
-  std::unique_ptr<trace::Session> TraceSession;
+  std::unique_ptr<trace::EventTracer> Tracer;
   if (!TraceFile.empty()) {
     std::error_code EC;
     TraceStream.emplace(TraceFile, /*ref*/ EC, llvm::sys::fs::F_RW);
@@ -124,10 +126,14 @@ int main(int argc, char *argv[]) {
       TraceFile.reset();
       llvm::errs() << "Error while opening trace file: " << EC.message();
     } else {
-      TraceSession = trace::Session::create(*TraceStream, PrettyPrint);
+      Tracer = trace::createJSONTracer(*TraceStream, PrettyPrint);
     }
   }
 
+  llvm::Optional<trace::Session> TracingSession;
+  if (Tracer)
+    TracingSession.emplace(*Tracer);
+
   llvm::raw_ostream &Outs = llvm::outs();
   llvm::raw_ostream &Logs = llvm::errs();
   JSONOutput Out(Outs, Logs,
diff --git a/unittests/clangd/TraceTests.cpp b/unittests/clangd/TraceTests.cpp
index 48ea8c11..0f734e97 100644
--- a/unittests/clangd/TraceTests.cpp
+++ b/unittests/clangd/TraceTests.cpp
@@ -7,6 +7,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "Context.h"
 #include "Trace.h"
 
 #include "llvm/ADT/DenseMap.h"
@@ -74,10 +75,11 @@ TEST(TraceTest, SmokeTest) {
   std::string JSON;
   {
     raw_string_ostream OS(JSON);
-    auto Session = trace::Session::create(OS);
+    auto JSONTracer = trace::createJSONTracer(OS);
+    trace::Session Session(*JSONTracer);
     {
-      trace::Span S("A");
-      trace::log("B");
+      trace::Span S(Context::empty(), "A");
+      trace::log(Context::empty(), "B");
     }
   }