Add tracing spans for input gate hold and wait

Author: shrima-cf

src/workerd/io/BUILD.bazel

@@ -271,6 +271,7 @@ wd_cc_library(
     hdrs = ["io-gate.h"],
     visibility = ["//visibility:public"],
     deps = [
+        ":trace",
         "@capnp-cpp//src/kj",
         "@capnp-cpp//src/kj:kj-async",
     ],

src/workerd/io/io-context.c++

@@ -338,7 +338,8 @@ IoContext::IncomingRequest::~IoContext_IncomingRequest() noexcept(false) {
 }
 
 InputGate::Lock IoContext::getInputLock() {
-  return KJ_ASSERT_NONNULL(currentInputLock, "no input lock available in this context").addRef();
+  return KJ_ASSERT_NONNULL(currentInputLock, "no input lock available in this context")
+      .addRef(getCurrentTraceSpan());
 }
 
 kj::Maybe<kj::Own<InputGate::CriticalSection>> IoContext::getCriticalSection() {

src/workerd/io/io-gate-test.c++

@@ -29,7 +29,7 @@ KJ_TEST("InputGate basics") {
     KJ_EXPECT(!promise2.poll(ws));
     KJ_EXPECT(!promise3.poll(ws));
 
-    auto lock2 = lock.addRef();
+    auto lock2 = lock.addRef(nullptr);
     { auto drop = kj::mv(lock); }
 
     KJ_EXPECT(!promise2.poll(ws));
@@ -299,7 +299,7 @@ KJ_TEST("InputGate critical section lock outlives critical section") {
   KJ_ASSERT(lock.isFor(gate));
 
   // Adding a ref and dropping it shouldn't cause trouble.
-  lock.addRef();
+  lock.addRef(nullptr);
 
   // The gate should still be locked
   auto waiter = gate.wait();

src/workerd/io/io-gate.c++

@@ -24,11 +24,15 @@ InputGate::~InputGate() noexcept {
   KJ_ASSERT(waiters.empty());
 }
 
-InputGate::Waiter::Waiter(
-    kj::PromiseFulfiller<Lock>& fulfiller, InputGate& gate, bool isChildWaiter)
+InputGate::Waiter::Waiter(kj::PromiseFulfiller<Lock>& fulfiller,
+    InputGate& gate,
+    bool isChildWaiter,
+    SpanParent parentSpan)
     : fulfiller(fulfiller),
       gate(&gate),
-      isChildWaiter(isChildWaiter) {
+      isChildWaiter(isChildWaiter),
+      waitSpan(parentSpan.newChild("input_gate_lock_wait"_kjc)),
+      lockSpanParent(kj::mv(parentSpan)) {
   gate.hooks.inputGateWaiterAdded();
   if (isChildWaiter) {
     gate.waitingChildren.add(*this);
@@ -47,13 +51,13 @@ InputGate::Waiter::~Waiter() noexcept(false) {
   }
 }
 
-kj::Promise<InputGate::Lock> InputGate::wait() {
+kj::Promise<InputGate::Lock> InputGate::wait(SpanParent parentSpan) {
   KJ_IF_SOME(e, brokenState.tryGet<kj::Exception>()) {
     return kj::cp(e);
   } else if (lockCount == 0) {
-    return Lock(*this);
+    return Lock(*this, kj::mv(parentSpan));
   } else {
-    return kj::newAdaptedPromise<Lock, Waiter>(*this, false);
+    return kj::newAdaptedPromise<Lock, Waiter>(*this, false, kj::mv(parentSpan));
   }
 }
 
@@ -65,10 +69,11 @@ kj::Promise<void> InputGate::onBroken() {
   }
 }
 
-InputGate::Lock::Lock(InputGate& gate)
+InputGate::Lock::Lock(InputGate& gate, SpanParent parentSpan)
     : gate(&gate),
       cs(gate.isCriticalSection ? kj::Maybe(kj::addRef(static_cast<CriticalSection&>(gate)))
-                                : kj::none) {
+                                : kj::none),
+      lockSpan(parentSpan.newChild("input_gate_lock_hold"_kjc)) {
   InputGate* gateToLock = &gate;
 
   KJ_IF_SOME(c, cs) {
@@ -107,11 +112,11 @@ void InputGate::releaseLock() {
     if (!waitingChildren.empty()) {
       auto& waiter = waitingChildren.front();
       waitingChildren.remove(waiter);
-      waiter.fulfiller.fulfill(Lock(*this));
+      waiter.fulfiller.fulfill(Lock(*this, kj::mv(waiter.lockSpanParent)));
     } else if (!waiters.empty()) {
       auto& waiter = waiters.front();
       waiters.remove(waiter);
-      waiter.fulfiller.fulfill(Lock(*this));
+      waiter.fulfiller.fulfill(Lock(*this, kj::mv(waiter.lockSpanParent)));
     }
   }
 }
@@ -167,7 +172,7 @@ InputGate::CriticalSection::~CriticalSection() noexcept(false) {
   }
 }
 
-kj::Promise<InputGate::Lock> InputGate::CriticalSection::wait() {
+kj::Promise<InputGate::Lock> InputGate::CriticalSection::wait(SpanParent parentSpan) {
   for (;;) {
     switch (state) {
       case NOT_STARTED: {
@@ -183,11 +188,12 @@ kj::Promise<InputGate::Lock> InputGate::CriticalSection::wait() {
         // Add ourselves to this parent's child waiter list.
         if (target.lockCount == 0) {
           state = RUNNING;
-          parentLock = Lock(target);
+          parentLock = Lock(target, parentSpan.addRef());
           continue;
         } else {
           try {
-            auto lock = co_await kj::newAdaptedPromise<Lock, Waiter>(target, true);
+            auto lock =
+                co_await kj::newAdaptedPromise<Lock, Waiter>(target, true, parentSpan.addRef());
             state = RUNNING;
             parentLock = kj::mv(lock);
             continue;
@@ -206,18 +212,18 @@ kj::Promise<InputGate::Lock> InputGate::CriticalSection::wait() {
         KJ_FAIL_REQUIRE("CriticalSection::wait() should be called once initially");
       case RUNNING:
         // CriticalSection is active, so defer to InputGate implementation.
-        co_return co_await InputGate::wait();
+        co_return co_await InputGate::wait(kj::mv(parentSpan));
       case REPARENTED:
         // Once the CriticalSection has declared itself done, then any straggler tasks it initiated
         // are adopted by the parent.
         // WARNING: Don't use parentAsInputGate() here as that'll bypass the override of wait() if
         //   the parent is a CriticalSection itself.
         KJ_SWITCH_ONEOF(parent) {
           KJ_CASE_ONEOF(p, InputGate*) {
-            co_return co_await p->wait();
+            co_return co_await p->wait(kj::mv(parentSpan));
           }
           KJ_CASE_ONEOF(c, kj::Own<CriticalSection>) {
-            co_return co_await c->wait();
+            co_return co_await c->wait(kj::mv(parentSpan));
           }
         }
         KJ_UNREACHABLE;

src/workerd/io/io-gate.h

@@ -20,6 +20,8 @@
 //   to the application are still being flushed to disk. If the flush fails, these messages will
 //   never be sent, so that the rest of the world cannot observe a prematurely-confirmed write.
 
+#include <workerd/io/trace.h>
+
 #include <kj/async.h>
 #include <kj/list.h>
 #include <kj/one-of.h>
@@ -62,7 +64,7 @@ class InputGate {
   class Lock {
    public:
     KJ_DISALLOW_COPY(Lock);
-    Lock(Lock&& other): gate(other.gate), cs(kj::mv(other.cs)) {
+    Lock(Lock&& other): gate(other.gate), cs(kj::mv(other.cs)), lockSpan(kj::mv(other.lockSpan)) {
       other.gate = nullptr;
     }
     ~Lock() noexcept(false) {
@@ -71,8 +73,8 @@ class InputGate {
 
     // Increments the lock's refcount, returning a duplicate `Lock`. All `Lock`s must be dropped
     // before the gate is unlocked.
-    Lock addRef() {
-      return Lock(*gate);
+    Lock addRef(SpanParent parentSpan) {
+      return Lock(*gate, kj::mv(parentSpan));
     }
 
     // Start a new critical section from this lock. After `wait()` has been called on the returned
@@ -99,12 +101,17 @@ class InputGate {
 
     kj::Maybe<kj::Own<CriticalSection>> cs;
 
-    Lock(InputGate& gate);
+    SpanBuilder lockSpan;
+
+    Lock(InputGate& gate, SpanParent parentSpan);
     friend class InputGate;
   };
 
   // Wait until there are no `Lock`s, then create a new one and return it.
-  kj::Promise<Lock> wait();
+  // If parentSpan is provided, child spans will be created to track:
+  // - Time spent waiting for the lock (if waiting is required)
+  // - Time spent holding the lock
+  kj::Promise<Lock> wait(SpanParent parentSpan = nullptr);
 
   // Rejects if and when calls to `wait()` become broken due to a failed critical section. The
   // actor should be shut down in this case. This promise never resolves, only rejects.
@@ -122,13 +129,21 @@ class InputGate {
   bool isCriticalSection = false;
 
   struct Waiter {
-    Waiter(kj::PromiseFulfiller<Lock>& fulfiller, InputGate& gate, bool isChildWaiter);
+    Waiter(kj::PromiseFulfiller<Lock>& fulfiller,
+        InputGate& gate,
+        bool isChildWaiter,
+        SpanParent parentSpan);
     ~Waiter() noexcept(false);
 
     kj::PromiseFulfiller<Lock>& fulfiller;
     InputGate* gate;
     bool isChildWaiter;
     kj::ListLink<Waiter> link;
+
+    // Span tracking how long we wait for the lock. Ends when Waiter is destroyed.
+    SpanBuilder waitSpan;
+    // Parent span to pass to Lock when it's created.
+    SpanParent lockSpanParent;
   };
 
   kj::List<Waiter, &Waiter::link> waiters;
@@ -173,7 +188,7 @@ class InputGate::CriticalSection: private InputGate, public kj::Refcounted {
   // The first call to wait() begins the CriticalSection. After that wait completes, until the
   // CriticalSection is done and dropped, no other locks will be allowed on this InputGate, except
   // locks requested by calling wait() on this CriticalSection -- or one of its children.
-  kj::Promise<Lock> wait();
+  kj::Promise<Lock> wait(SpanParent parentSpan = nullptr);
 
   // Call when the critical section has completed successfully. If this is not called before the
   // CriticalSection is dropped, then failed() is called implicitly.

src/workerd/io/worker.c++

@@ -3793,7 +3793,7 @@ kj::Promise<void> Worker::Actor::ensureConstructedImpl(IoContext& context, Actor
       handler.missingSuperclass = info.missingSuperclass;
 
       impl->classInstance = kj::mv(handler);
-    }, inputLock.addRef());
+    }, inputLock.addRef(context.getCurrentTraceSpan()));
     // We addRef() the inputLock above rather than kj::mv() it so that the lock remains held
     // through the catch block below, if an exception is thrown. This is important since we
     // MUST update `impl->classInstance` to something other than `Initializing` before we