[clangd] Fix tracing now that spans lifetimes can overlap on a thread.

Summary:
The chrome trace viewer requires events within a thread to strictly nest.
So we need to record the lifetime of the Span objects, not the contexts.

But we still want to show the relationship between spans where a context crosses
threads, so do this with flow events (i.e. arrows).

Before: https://photos.app.goo.gl/q4Dd9u9xtelaXk1v1
After: https://photos.app.goo.gl/5RNLmAMLZR3unvY83

(This could stand some further improvement, in particular I think we want a
container span whenever we schedule work on a thread. But that's another patch)

Reviewers: ioeric

Subscribers: klimek, ilya-biryukov, jkorous-apple, cfe-commits

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

llvm-svn: 325220

clang-tools-extra/clangd/Trace.cpp

@@ -16,6 +16,7 @@
 #include "llvm/Support/FormatProviders.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Threading.h"
+#include <atomic>
 #include <mutex>
 
 namespace clang {
@@ -46,23 +47,117 @@ public:
     Out.flush();
   }
 
+  // We stash a Span object in the context. It will record the start/end,
+  // and this also allows us to look up the parent Span's information.
   Context beginSpan(llvm::StringRef Name, json::obj *Args) override {
-    jsonEvent("B", json::obj{{"name", Name}});
-    return Context::current().derive(make_scope_exit([this, Args] {
-      jsonEvent("E", json::obj{{"args", std::move(*Args)}});
-    }));
+    return Context::current().derive(SpanKey,
+                                     make_unique<JSONSpan>(this, Name, Args));
+  }
+
+  // Trace viewer requires each thread to properly stack events.
+  // So we need to mark only duration that the span was active on the thread.
+  // (Hopefully any off-thread activity will be connected by a flow event).
+  // Record the end time here, but don't write the event: Args aren't ready yet.
+  void endSpan() override {
+    Context::current().getExisting(SpanKey)->markEnded();
   }
 
   void instant(llvm::StringRef Name, json::obj &&Args) override {
+    captureThreadMetadata();
     jsonEvent("i", json::obj{{"name", Name}, {"args", std::move(Args)}});
   }
 
   // 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 jsonEvent(StringRef Phase, json::obj &&Contents) {
+  void jsonEvent(StringRef Phase, json::obj &&Contents,
+                 uint64_t TID = get_threadid(),
+                 double Timestamp = 0) {
+    Contents["ts"] = Timestamp ? Timestamp : timestamp();
+    Contents["tid"] = TID;
+    std::lock_guard<std::mutex> Lock(Mu);
+    rawEvent(Phase, std::move(Contents));
+  }
+
+private:
+  class JSONSpan {
+  public:
+    JSONSpan(JSONTracer *Tracer, llvm::StringRef Name, json::obj *Args)
+        : StartTime(Tracer->timestamp()), EndTime(0), Name(Name),
+          TID(get_threadid()), Tracer(Tracer), Args(Args) {
+      // ~JSONSpan() may run in a different thread, so we need to capture now.
+      Tracer->captureThreadMetadata();
+
+      // We don't record begin events here (and end events in the destructor)
+      // because B/E pairs have to appear in the right order, which is awkward.
+      // Instead we send the complete (X) event in the destructor.
+
+      // If our parent was on a different thread, add an arrow to this span.
+      auto *Parent = Context::current().get(SpanKey);
+      if (Parent && *Parent && (*Parent)->TID != TID) {
+        // If the parent span ended already, then show this as "following" it.
+        // Otherwise show us as "parallel".
+        double OriginTime = (*Parent)->EndTime;
+        if (!OriginTime)
+          OriginTime = (*Parent)->StartTime;
+
+        auto FlowID = nextID();
+        Tracer->jsonEvent("s",
+                          json::obj{{"id", FlowID},
+                                    {"name", "Context crosses threads"},
+                                    {"cat", "dummy"}},
+                          (*Parent)->TID, (*Parent)->StartTime);
+        Tracer->jsonEvent("f",
+                          json::obj{{"id", FlowID},
+                                    {"bp", "e"},
+                                    {"name", "Context crosses threads"},
+                                    {"cat", "dummy"}},
+                          TID);
+      }
+    }
+
+    ~JSONSpan() {
+      // Finally, record the event (ending at EndTime, not timestamp())!
+      Tracer->jsonEvent("X",
+                        json::obj{{"name", std::move(Name)},
+                                  {"args", std::move(*Args)},
+                                  {"dur", EndTime - StartTime}},
+                        TID, StartTime);
+    }
+
+    // May be called by any thread.
+    void markEnded() {
+      EndTime = Tracer->timestamp();
+    }
+
+  private:
+    static uint64_t nextID() {
+      static std::atomic<uint64_t> Next = {0};
+      return Next++;
+    }
+
+    double StartTime;
+    std::atomic<double> EndTime; // Filled in by markEnded().
+    std::string Name;
+    uint64_t TID;
+    JSONTracer *Tracer;
+    json::obj *Args;
+  };
+  static Key<std::unique_ptr<JSONSpan>> SpanKey;
+
+  // Record an event. ph and pid are set.
+  // Contents must be a list of the other JSON key/values.
+  void rawEvent(StringRef Phase, json::obj &&Event) /*REQUIRES(Mu)*/ {
+    // PID 0 represents the clangd process.
+    Event["pid"] = 0;
+    Event["ph"] = Phase;
+    Out << Sep << formatv(JSONFormat, json::Expr(std::move(Event)));
+    Sep = ",\n";
+  }
+
+  // If we haven't already, emit metadata describing this thread.
+  void captureThreadMetadata() {
     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);
@@ -74,20 +169,6 @@ public:
                       });
       }
     }
-    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.
-  void rawEvent(StringRef Phase, json::obj &&Event) /*REQUIRES(Mu)*/ {
-    // PID 0 represents the clangd process.
-    Event["pid"] = 0;
-    Event["ph"] = Phase;
-    Out << Sep << formatv(JSONFormat, json::Expr(std::move(Event)));
-    Sep = ",\n";
   }
 
   double timestamp() {
@@ -103,6 +184,8 @@ private:
   const char *JSONFormat;
 };
 
+Key<std::unique_ptr<JSONTracer::JSONSpan>> JSONTracer::SpanKey;
+
 EventTracer *T = nullptr;
 } // namespace
 
@@ -139,6 +222,11 @@ Span::Span(llvm::StringRef Name)
     : Args(T ? new json::obj() : nullptr),
       RestoreCtx(makeSpanContext(Name, Args)) {}
 
+Span::~Span() {
+  if (T)
+    T->endSpan();
+}
+
 } // namespace trace
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/Trace.h

@@ -40,6 +40,12 @@ public:
   /// whose destructor records the end of the event.
   /// The args are *Args, only complete when the event ends.
   virtual Context beginSpan(llvm::StringRef Name, json::obj *Args) = 0;
+  // Called when a Span is destroyed (it may still be active on other threads).
+  // beginSpan() and endSpan() will always form a proper stack on each thread.
+  // The Context returned by beginSpan is active, but Args is not ready.
+  // Tracers should not override this unless they need to observe strict
+  // per-thread nesting. Instead they should observe context destruction.
+  virtual void endSpan() {};
 
   /// Called for instant events.
   virtual void instant(llvm::StringRef Name, json::obj &&Args) = 0;
@@ -77,6 +83,7 @@ void log(const llvm::Twine &Name);
 class Span {
 public:
   Span(llvm::StringRef Name);
+  ~Span();
 
   /// Mutable metadata, if this span is interested.
   /// Prefer to use SPAN_ATTACH rather than accessing this directly.

clang-tools-extra/test/clangd/trace.test

@@ -2,18 +2,22 @@
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
 ---
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{"uri":"test:///foo.c","languageId":"c","version":1,"text":"void main() {}"}}}
-#      CHECK: {"displayTimeUnit":"ns","traceEvents":[
-# Start opening the doc.
-#      CHECK: "name": "textDocument/didOpen"
-#      CHECK: "ph": "B"
-# 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"
+# These assertions are a bit loose, to avoid brittleness.
+# CHECK: {"displayTimeUnit":"ns","traceEvents":[
+# CHECK: {
+# CHECK:   "args": {
+# CHECK:     "File": "{{.*(/|\\)}}foo.c"
+# CHECK:   },
+# CHECK:   "name": "Preamble",
+# CHECK:   "ph": "X",
+# CHECK: }
+# CHECK: {
+# CHECK:   "args": {
+# CHECK:     "File": "{{.*(/|\\)}}foo.c"
+# CHECK:   },
+# CHECK:   "name": "Build",
+# CHECK:   "ph": "X",
+# CHECK: }
+# CHECK: },
 ---
 {"jsonrpc":"2.0","id":5,"method":"shutdown"}

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

@@ -114,12 +114,10 @@ TEST(TraceTest, SmokeTest) {
     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", "Log"}}));
   ASSERT_NE(++Event, Events->end()) << "Expected span end";
-  EXPECT_TRUE(VerifyObject(*Event, {{"ph", "E"}}));
+  EXPECT_TRUE(VerifyObject(*Event, {{"ph", "X"}, {"name", "A"}}));
   ASSERT_EQ(++Event, Events->end());
   ASSERT_EQ(++Prop, Root->end());
 }