[clangd] Tracing improvements

[sammccall]

Summary

[clangd] Tracing improvements

• Compose JSON using JSONExpr
• Allow attaching metadata to spans (and avoid it if tracing is off)
• Attach IDs and responses of JSON RPCs to their spans
• When -pretty is passed, also prettyprint the trace

The downside is that large responses make the trace viewer sluggish.
We should make our responses less huge :-) Or fix trace viewer.

[ilya-biryukov]

clangd/JSONRPCDispatcher.h

61

NIT: maybe replace new with llvm::make_unique?

78

Why do we need unique_ptr? Are Spans non-movable?

clangd/Trace.cpp

50

Any reason why we use rval-ref instead of accepting by-value?

[ilya-biryukov]

Oops, accidentally accepted the revision.
There are no major issues, but still wanted to learn the reasons we use rval-refs and allocate span dynamically.

[sammccall]

Address review comment.

clangd/JSONRPCDispatcher.h

78

Spans aren't movable, they have an explicitly declared destructor. (The copy constructor is only deprecated by this condition, but Span has a unique_ptr member).

We could make them movable, though it's not absolutely trivial (we need an extra bool to indicate that this is moved-from and the destructor should be a no-op).

I think wrapping them in a unique_ptr here is slightly simpler than implementing the right move semantics by hand, but LMK what you think.

clangd/Trace.cpp

50

Two reasons I prefer this for json::expr/obj/arr:

• major: you almost always want to pass a new literal, or move an existing one. Passing a const Expr & is usually a mistake, and taking Expr&& makes it an error.
• minor: expr/obj/arr aren't trivially cheap to move. We forward these around internally, taking && at every level means we only pay for the move constructor once, pass-by-value pays on every call.

[ilya-biryukov]

LGTM

clangd/JSONRPCDispatcher.h

78

We could make them movable, though it's not absolutely trivial (we need an extra bool to indicate that this is moved-from and the destructor should be a no-op).

I kinda hate this part when dealing with movable objects, too. I still do that, mostly to avoid heap allocs, but we're not on a hot path, so it's fine both ways.
We could use llvm::Optional instead of unique_ptr to get rid of the heap alloc too.

clangd/Trace.cpp

50

Makes sense, thanks for expanding on this. I'll make sure to pass around by r-value ref too when dealing with JSON objects.

[ilya-biryukov]

clangd/JSONRPCDispatcher.h

78

Oh, sorry, llvm::Optional surely doesn't help here, it would still be non-movable. Forget what I said.

clangd/ClangdUnit.cpp

auto FinishRebuild = [OldPreamble, VFS, RequestRebuildCounter, PCHs,
     auto DoRebuildPreamble = [&]() -> std::shared_ptr<const PreambleData> {
       auto Bounds =
           ComputePreambleBounds(*CI->getLangOpts(), ContentsBuffer.get(), 0);
       if (OldPreamble && OldPreamble->Preamble.CanReuse(
                              *CI, ContentsBuffer.get(), Bounds, VFS.get())) {
         return OldPreamble;
       }
 
-      trace::Span Tracer(llvm::Twine("Preamble: ") + That->FileName);
+      trace::Span Tracer("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(
           *CI, ContentsBuffer.get(), Bounds, *PreambleDiagsEngine, VFS, PCHs,
if (NewPreamble) {
       SerializedPreambleDecls = NewPreamble->TopLevelDeclIDs;
       Diagnostics.insert(Diagnostics.begin(), NewPreamble->Diags.begin(),
                          NewPreamble->Diags.end());
     }
 
     // Compute updated AST.
     llvm::Optional<ParsedAST> NewAST;
     {
-      trace::Span Tracer(llvm::Twine("Build: ") + That->FileName);
+      trace::Span Tracer("Build");
+      SPAN_ATTACH(Tracer, "File", That->FileName);
       NewAST = ParsedAST::Build(
           std::move(CI), PreambleForAST, SerializedPreambleDecls,
           std::move(ContentsBuffer), PCHs, VFS, That->Logger);
     }
 
     if (NewAST) {
       Diagnostics.insert(Diagnostics.end(), NewAST->getDiagnostics().begin(),
                          NewAST->getDiagnostics().end());

clangd/JSONRPCDispatcher.h

 //===--- JSONRPCDispatcher.h - Main JSON parser entry point -----*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_JSONRPCDISPATCHER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_JSONRPCDISPATCHER_H
 
 #include "JSONExpr.h"
 #include "Logger.h"
 #include "Protocol.h"
+#include "Trace.h"
 #include "clang/Basic/LLVM.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/YAMLParser.h"
 #include <iosfwd>
 #include <mutex>
 
 namespace clang {
private:
   bool Pretty;
 
   std::mutex StreamMutex;
 };
 
 /// Context object passed to handlers to allow replies.
 class RequestContext {
 public:
-  RequestContext(JSONOutput &Out, llvm::Optional<json::Expr> ID)
-      : Out(Out), ID(std::move(ID)) {}
+  RequestContext(JSONOutput &Out, StringRef Method,
+                 llvm::Optional<json::Expr> ID)
+      : Out(Out), ID(std::move(ID)), Tracer(new trace::Span(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; }
+
 private:
   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 {
 public:
   // A handler responds to requests for a particular method name.
   using Handler =

clangd/JSONRPCDispatcher.cpp

void JSONOutput::mirrorInput(const Twine &Message) {
   InputMirror->flush();
 }
 
 void RequestContext::reply(json::Expr &&Result) {
   if (!ID) {
     Out.log("Attempted to reply to a notification!\n");
     return;
   }
+  SPAN_ATTACH(tracer(), "Reply", Result);
   Out.writeMessage(json::obj{
       {"jsonrpc", "2.0"},
       {"id", *ID},
       {"result", std::move(Result)},
   });
 }
 
 void RequestContext::replyError(ErrorCode code, const llvm::StringRef &Message) {
   Out.log("Error " + Twine(static_cast<int>(code)) + ": " + Message + "\n");
+  SPAN_ATTACH(tracer(), "Error",
+              (json::obj{{"code", static_cast<int>(code)},
+                         {"message", Message.str()}}));
   if (ID) {
     Out.writeMessage(json::obj{
         {"jsonrpc", "2.0"},
         {"id", *ID},
         {"error", json::obj{{"code", static_cast<int>(code)}, {"message", Message}}},
     });
   }
 }
 
 void RequestContext::call(StringRef Method, json::Expr &&Params) {
   // FIXME: Generate/Increment IDs for every request so that we can get proper
   // replies once we need to.
+  SPAN_ATTACH(tracer(), "Call",
+              (json::obj{{"method", Method.str()}, {"params", Params}}));
   Out.writeMessage(json::obj{
       {"jsonrpc", "2.0"},
       {"id", 1},
       {"method", Method},
       {"params", std::move(Params)},
   });
 }
 
 void JSONRPCDispatcher::registerHandler(StringRef Method, Handler H) {
   assert(!Handlers.count(Method) && "Handler already registered!");
   Handlers[Method] = std::move(H);
 }
 
 static void
 callHandler(const llvm::StringMap<JSONRPCDispatcher::Handler> &Handlers,
             llvm::yaml::ScalarNode *Method, llvm::Optional<json::Expr> ID,
             llvm::yaml::MappingNode *Params,
             const JSONRPCDispatcher::Handler &UnknownHandler, JSONOutput &Out) {
   llvm::SmallString<64> MethodStorage;
   llvm::StringRef MethodStr = Method->getValue(MethodStorage);
   auto I = Handlers.find(MethodStr);
   auto &Handler = I != Handlers.end() ? I->second : UnknownHandler;
-  trace::Span Tracer(MethodStr);
-  Handler(RequestContext(Out, std::move(ID)), Params);
+  Handler(RequestContext(Out, MethodStr, std::move(ID)), Params);
 }
 
 bool JSONRPCDispatcher::call(StringRef Content, JSONOutput &Out) const {
   llvm::SourceMgr SM;
   llvm::yaml::Stream YAMLStream(Content, SM);
 
   auto Doc = YAMLStream.begin();
   if (Doc == YAMLStream.end())
if (ContentLength > 1 << 30) { // 1024M
               " bytes.\n");
       continue;
     }
 
     if (ContentLength > 0) {
       std::vector<char> JSON(ContentLength + 1, '\0');
       llvm::StringRef JSONRef;
       {
-        trace::Span Tracer("Reading request");
         // Now read the JSON. Insert a trailing null byte as required by the
         // YAML parser.
         In.read(JSON.data(), ContentLength);
         Out.mirrorInput(StringRef(JSON.data(), In.gcount()));
 
         // If the stream is aborted before we read ContentLength bytes, In
         // will have eofbit and failbit set.
         if (!In) {

clangd/Trace.h

 //
 // 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 "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 {
 public:
   // Starts a sessions capturing trace events and writing Trace Event JSON.
-  static std::unique_ptr<Session> create(llvm::raw_ostream &OS);
+  static std::unique_ptr<Session> create(llvm::raw_ostream &OS,
+                                         bool Pretty = false);
   ~Session();
 
 private:
   Session() = default;
 };
 
 // Records a single instant event, associated with the current thread.
 void log(const llvm::Twine &Name);
 
 // 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.
 class Span {
 public:
-  Span(const llvm::Twine &Name);
+  Span(std::string Name);
   ~Span();
 
+  // 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;
 };
 
+#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 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 {
 public:
-  Tracer(raw_ostream &Out)
-      : Out(Out), Sep(""), Start(std::chrono::system_clock::now()) {
+  Tracer(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", R"("name": "process_name", "args":{"name":"clangd"})");
+    rawEvent("M", json::obj{
+                      {"name", "process_name"},
+                      {"args", json::obj{{"name", "clangd"}}},
+                  });
   }
 
   ~Tracer() {
     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.
-  template <typename T> void event(StringRef Phase, const T &Contents) {
+  void event(StringRef Phase, json::obj &&Contents) {
     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",
-            formatv(
-                R"("tid": {0}, "name": "thread_name", "args":{"name":"{1}"})",
-                TID, StringRef(&Name[0], Name.size())));
+        rawEvent("M", json::obj{
+                          {"tid", TID},
+                          {"name", "thread_name"},
+                          {"args", json::obj{{"name", Name}}},
+                      });
       }
     }
-    rawEvent(Phase, formatv(R"("ts":{0}, "tid":{1}, {2})", timestamp(), TID,
-                            Contents));
+    Contents["ts"] = timestamp();
+    Contents["tid"] = TID;
+    rawEvent(Phase, std::move(Contents));
   }
 
 private:
   // Record an event. ph and pid are set.
   // Contents must be a list of the other JSON key/values.
-  template <typename T>
-  void rawEvent(StringRef Phase, const T &Contents) /*REQUIRES(Mu)*/ {
+  void rawEvent(StringRef Phase, json::obj &&Event) /*REQUIRES(Mu)*/ {
     // PID 0 represents the clangd process.
-    Out << Sep << R"({"pid":0, "ph":")" << Phase << "\", " << Contents << "}";
+    Event["pid"] = 0;
+    Event["ph"] = Phase;
+    Out << Sep << formatv(JSONFormat, json::Expr(std::move(Event)));
     Sep = ",\n";
   }
 
   double timestamp() {
     using namespace std::chrono;
-    return duration<double, std::milli>(system_clock::now() - Start).count();
+    return duration<double, std::micro>(system_clock::now() - Start).count();
   }
 
   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) {
+std::unique_ptr<Session> Session::create(raw_ostream &OS, bool Pretty) {
   assert(!T && "A session is already active");
-  T = new Tracer(OS);
+  T = new Tracer(OS, Pretty);
   return std::unique_ptr<Session>(new Session());
 }
 
 Session::~Session() {
   delete T;
   T = nullptr;
 }
 
 void log(const Twine &Message) {
   if (!T)
     return;
-  T->event("i", formatv(R"("name":"{0}")", yaml::escape(Message.str())));
+  T->event("i", json::obj{
+                    {"name", "Log"},
+                    {"args", json::obj{{"Message", Message.str()}}},
+                });
 }
 
-Span::Span(const Twine &Text) {
+Span::Span(std::string Name) {
   if (!T)
     return;
-  T->event("B", formatv(R"("name":"{0}")", yaml::escape(Text.str())));
+  T->event("B", json::obj{{"name", std::move(Name)}});
+  Args = llvm::make_unique<json::obj>();
 }
 
 Span::~Span() {
   if (!T)
     return;
-  T->event("E", R"("_":0)" /* Dummy property to ensure valid JSON */);
+  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

int main(int argc, char *argv[]) {
   std::unique_ptr<trace::Session> 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);
+      TraceSession = trace::Session::create(*TraceStream, PrettyPrint);
     }
   }
 
   llvm::raw_ostream &Outs = llvm::outs();
   llvm::raw_ostream &Logs = llvm::errs();
   JSONOutput Out(Outs, Logs,
                  InputMirrorStream ? InputMirrorStream.getPointer() : nullptr,
                  PrettyPrint);

test/clangd/trace.test

-# RUN: clangd -run-synchronously -trace %t < %s && FileCheck %s < %t
+# RUN: clangd -pretty -run-synchronously -trace %t < %s && FileCheck %s < %t
 # It is absolutely vital that this file has CRLF line endings.
 #
 Content-Length: 125
 
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
 #
 Content-Length: 152
 
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{"uri":"file:///foo.c","languageId":"c","version":1,"text":"void main() {}"}}}
-# CHECK: "textDocument/didOpen"
-# CHECK: "Preamble: /foo.c"
-# CHECK: "Build: /foo.c"
+#      CHECK: {"displayTimeUnit":"ns","traceEvents":[
+# Start opening the doc.
+#      CHECK: "name": "textDocument/didOpen"
+#      CHECK: "ph": "E"
+# Start building the preamble.
+#      CHECK: "name": "Preamble"
+#      CHECK: },
+# Finish building the preamble, with filename.
+#      CHECK: "File": "/foo.c"
+# CHECK-NEXT: },
+# CHECK-NEXT: "ph": "E"
+# Start building the file.
+#      CHECK: "name": "Build"
 Content-Length: 44
 
 {"jsonrpc":"2.0","id":5,"method":"shutdown"}

unittests/clangd/TraceTests.cpp

TEST(TraceTest, SmokeTest) {
   EXPECT_TRUE(VerifyObject(*Event, {{"ph", "M"}, {"name", "process_name"}}));
   if (ThreadsHaveNames) {
     ASSERT_NE(++Event, Events->end()) << "Expected thread name";
     EXPECT_TRUE(VerifyObject(*Event, {{"ph", "M"}, {"name", "thread_name"}}));
   }
   ASSERT_NE(++Event, Events->end()) << "Expected span start";
   EXPECT_TRUE(VerifyObject(*Event, {{"ph", "B"}, {"name", "A"}}));
   ASSERT_NE(++Event, Events->end()) << "Expected log message";
-  EXPECT_TRUE(VerifyObject(*Event, {{"ph", "i"}, {"name", "B"}}));
+  EXPECT_TRUE(VerifyObject(*Event, {{"ph", "i"}, {"name", "Log"}}));
   ASSERT_NE(++Event, Events->end()) << "Expected span end";
   EXPECT_TRUE(VerifyObject(*Event, {{"ph", "E"}}));
   ASSERT_EQ(++Event, Events->end());
   ASSERT_EQ(++Prop, Root->end());
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang