Merge pull request #5511 from cloudflare/felix/111125-stw-simplify

[o11y][nfc] Deduplicate user tracing observer with internal version

Author: fhanau

src/workerd/io/tracer.c++

@@ -504,4 +504,17 @@ void WorkerTracer::setJsRpcInfo(const tracing::InvocationSpanContext& context,
   }
 }
 
+kj::Own<SpanObserver> UserSpanObserver::newChild() {
+  return kj::refcounted<UserSpanObserver>(kj::addRef(*submitter), spanId);
+}
+
+void UserSpanObserver::report(const Span& span) {
+  submitter->submitSpan(spanId, parentSpanId, span);
+}
+
+// Provide I/O time to the tracing system for user spans.
+kj::Date UserSpanObserver::getTime() {
+  return IoContext::current().now();
+}
+
 }  // namespace workerd

src/workerd/io/tracer.h

@@ -178,4 +178,36 @@ class WorkerTracer final: public BaseTracer {
 
   kj::Maybe<kj::Own<tracing::TailStreamWriter>> maybeTailStreamWriter;
 };
+
+class SpanSubmitter: public kj::Refcounted {
+ public:
+  virtual void submitSpan(tracing::SpanId context, tracing::SpanId spanId, const Span& span) = 0;
+  virtual tracing::SpanId makeSpanId() = 0;
+};
+
+// The user tracing observer
+class UserSpanObserver final: public SpanObserver {
+ public:
+  // constructor for top-level observer
+  UserSpanObserver(kj::Own<SpanSubmitter> submitter)
+      : submitter(kj::mv(submitter)),
+        spanId(tracing::SpanId::nullId),
+        parentSpanId(tracing::SpanId::nullId) {}
+  // constructor for subsequent observers attached to a span
+  UserSpanObserver(kj::Own<SpanSubmitter> submitter, tracing::SpanId parentSpanId)
+      : submitter(kj::mv(submitter)),
+        spanId(this->submitter->makeSpanId()),
+        parentSpanId(parentSpanId) {}
+  KJ_DISALLOW_COPY(UserSpanObserver);
+
+  kj::Own<SpanObserver> newChild() override;
+  void report(const Span& span) override;
+  kj::Date getTime() override;
+
+ private:
+  kj::Own<SpanSubmitter> submitter;
+  tracing::SpanId spanId;
+  tracing::SpanId parentSpanId;
+};
+
 }  // namespace workerd

src/workerd/server/server.c++

@@ -1699,10 +1699,10 @@ class RequestObserverWithTracer final: public RequestObserver, public WorkerInte
   kj::uint fetchStatus = 0;
 };
 
-class SpanSubmitter final: public kj::Refcounted {
+class SequentialSpanSubmitter final: public SpanSubmitter {
  public:
-  SpanSubmitter(kj::Own<WorkerTracer> workerTracer): workerTracer(kj::mv(workerTracer)) {}
-  void submitSpan(tracing::SpanId spanId, tracing::SpanId parentSpanId, const Span& span) {
+  SequentialSpanSubmitter(kj::Own<WorkerTracer> workerTracer): workerTracer(kj::mv(workerTracer)) {}
+  void submitSpan(tracing::SpanId spanId, tracing::SpanId parentSpanId, const Span& span) override {
     // We largely recreate the span here which feels inefficient, but is hard to avoid given the
     // mismatch between the Span type and the full span information required for OTel.
     tracing::CompleteSpan span2(
@@ -1718,45 +1718,16 @@ class SpanSubmitter final: public kj::Refcounted {
     workerTracer->addSpan(kj::mv(span2));
   }
 
-  tracing::SpanId makeSpanId() {
-    return tracing::SpanId(predictableSpanId++);
+  tracing::SpanId makeSpanId() override {
+    return tracing::SpanId(nextSpanId++);
   }
-  KJ_DISALLOW_COPY_AND_MOVE(SpanSubmitter);
+  KJ_DISALLOW_COPY_AND_MOVE(SequentialSpanSubmitter);
 
  private:
-  uint64_t predictableSpanId = 0;
+  uint64_t nextSpanId = 1;
   kj::Own<WorkerTracer> workerTracer;
 };
 
-class WorkerTracerSpanObserver: public SpanObserver {
- public:
-  WorkerTracerSpanObserver(
-      kj::Own<SpanSubmitter> spanSubmitter, tracing::SpanId parentSpanId = tracing::SpanId::nullId)
-      : spanSubmitter(kj::mv(spanSubmitter)),
-        spanId(this->spanSubmitter->makeSpanId()),
-        parentSpanId(parentSpanId) {}
-
-  KJ_DISALLOW_COPY_AND_MOVE(WorkerTracerSpanObserver);
-
-  [[nodiscard]] kj::Own<SpanObserver> newChild() override {
-    return kj::refcounted<WorkerTracerSpanObserver>(kj::addRef(*spanSubmitter), spanId);
-  }
-
-  void report(const Span& span) override {
-    spanSubmitter->submitSpan(spanId, parentSpanId, span);
-  }
-
-  // Provide I/O time to the tracing system for user spans.
-  kj::Date getTime() override {
-    return IoContext::current().now();
-  }
-
- private:
-  kj::Own<SpanSubmitter> spanSubmitter;
-  tracing::SpanId spanId;
-  tracing::SpanId parentSpanId;
-};
-
 // IsolateLimitEnforcer that enforces no limits.
 class NullIsolateLimitEnforcer final: public IsolateLimitEnforcer {
  public:
@@ -2201,8 +2172,8 @@ class Server::WorkerService final: public Service,
 
     KJ_IF_SOME(w, workerTracer) {
       w->setMakeUserRequestSpanFunc([&w = *w]() {
-        return SpanParent(
-            kj::refcounted<WorkerTracerSpanObserver>(kj::refcounted<SpanSubmitter>(kj::addRef(w))));
+        return SpanParent(kj::refcounted<UserSpanObserver>(
+            kj::refcounted<SequentialSpanSubmitter>(kj::addRef(w))));
       });
     }
     kj::Own<RequestObserver> observer =