Index: clangd/JSONRPCDispatcher.cpp =================================================================== --- clangd/JSONRPCDispatcher.cpp +++ clangd/JSONRPCDispatcher.cpp @@ -20,7 +20,6 @@ using namespace clangd; namespace { -static Key> RequestSpan; static Key RequestID; static Key RequestOut; } // namespace @@ -66,8 +65,8 @@ 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{ @@ -80,10 +79,11 @@ void clangd::replyError(const Context &Ctx, ErrorCode code, const llvm::StringRef &Message) { log(Ctx, "Error " + Twine(static_cast(code)) + ": " + Message); - if (auto *Span = Ctx.get(RequestSpan)) - SPAN_ATTACH(**Span, "Error", - (json::obj{{"code", static_cast(code)}, - {"message", Message.str()}})); + + trace::Span Span(Ctx, "clangd::replyError"); + SPAN_ATTACH(Span, "Error", + (json::obj{{"code", static_cast(code)}, + {"message", Message.str()}})); if (auto ID = Ctx.get(RequestID)) { Ctx.getExisting(RequestOut) @@ -99,9 +99,11 @@ 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", + (json::obj{{"method", Method.str()}, {"params", Params}})); + Ctx.getExisting(RequestOut) ->writeMessage(json::obj{ {"jsonrpc", "2.0"}, @@ -137,19 +139,17 @@ 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(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; Index: clangd/Trace.h =================================================================== --- clangd/Trace.h +++ clangd/Trace.h @@ -28,21 +28,29 @@ 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; - 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, - llvm::StringRef Name) = 0; + virtual std::unique_ptr beginSpan(const Context &Ctx, + llvm::StringRef Name) = 0; /// Called for instant events. virtual void instant(const Context &Ctx, llvm::StringRef Name, @@ -78,15 +86,29 @@ 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 Args; - EventTracer::EndEventCallback Callback; + Context EventCtx; }; #define SPAN_ATTACH(S, Name, Expr) \ Index: clangd/Trace.cpp =================================================================== --- clangd/Trace.cpp +++ clangd/Trace.cpp @@ -24,6 +24,25 @@ // 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 ArgsObj; + JSONTracer &Tracer; + }; + public: JSONTracer(raw_ostream &Out, bool Pretty) : Out(Out), Sep(""), Start(std::chrono::system_clock::now()), @@ -43,14 +62,12 @@ Out.flush(); } - EndEventCallback beginSpan(const Context &Ctx, - llvm::StringRef Name) override { + std::unique_ptr 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(*this); } void instant(const Context &Ctx, llvm::StringRef Name, @@ -125,23 +142,54 @@ 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 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 Event; +}; + +static Key EnvelopeKey; +} // namespace + Span::Span(const Context &Ctx, llvm::StringRef Name) { if (!T) return; - Callback = T->beginSpan(Ctx, Name); - if (!Callback) + std::unique_ptr Ev = T->beginSpan(Ctx, Name); + if (!Ev) return; + EventCtx = Ctx.derive(EnvelopeKey, EventEnvelope(std::move(Ev))); Args = llvm::make_unique(); } +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