Use TraceParentContext, simplifying code using SubrequestMetadata

cloudflare · Sep 16, 2024 · b375e4f · b375e4f
1 parent 5df42cc
commit b375e4f
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 7 deletions.
diff --git a/src/workerd/api/actor.c++ b/src/workerd/api/actor.c++
@@ -35,7 +35,7 @@ public:
       }
 
       return KJ_REQUIRE_NONNULL(actorChannel)
-          ->startRequest({.cfBlobJson = kj::mv(cfStr), .parentSpan = spans.span});
+          ->startRequest({.cfBlobJson = kj::mv(cfStr), .spans = spans});
     },
         {.inHouse = true,
           .wrapMetrics = true,
@@ -77,7 +77,7 @@ public:
       }
 
       return KJ_REQUIRE_NONNULL(actorChannel)
-          ->startRequest({.cfBlobJson = kj::mv(cfStr), .parentSpan = spans.span});
+          ->startRequest({.cfBlobJson = kj::mv(cfStr), .spans = spans});
     },
         {.inHouse = true,
           .wrapMetrics = true,

diff --git a/src/workerd/io/io-channels.h b/src/workerd/io/io-channels.h
@@ -88,8 +88,7 @@ class IoChannelFactory {
     kj::Maybe<kj::String> cfBlobJson;
 
     // Specifies the parent span for the subrequest for tracing purposes.
-    SpanParent parentSpan = nullptr;
-    LimeSpanParent limeParentSpan = nullptr;
+    TraceParentContext spans = TraceParentContext(nullptr, nullptr);
 
     // Serialized JSON value to pass in ew_compat field of control header to FL. If this subrequest
     // does not go directly to FL, this value is ignored. Flags marked with `$neededByFl` in

diff --git a/src/workerd/io/io-context.c++ b/src/workerd/io/io-context.c++
@@ -852,8 +852,7 @@ kj::Own<WorkerInterface> IoContext::getSubrequestChannelImpl(uint channel,
     IoChannelFactory& channelFactory) {
   IoChannelFactory::SubrequestMetadata metadata{
     .cfBlobJson = kj::mv(cfBlobJson),
-    .parentSpan = spans.span,
-    .limeParentSpan = spans.limeSpan,
+    .spans = spans,
     .featureFlagsForFl = worker->getIsolate().getFeatureFlagsForFl(),
   };
 

diff --git a/src/workerd/io/trace.h b/src/workerd/io/trace.h
@@ -753,4 +753,19 @@ struct TraceContext {
   LimeSpanBuilder limeSpan;
 };
 
+// TraceContext variant tracking span parents instead. This is useful when we need for code
+// interacting with IoChannelFactory::SubrequestMetadata.
+struct TraceParentContext {
+  TraceParentContext(TraceContext& spans): span(spans.span), limeSpan(spans.limeSpan) {}
+  TraceParentContext(SpanParent span, LimeSpanParent limeSpan)
+      : span(kj::mv(span)),
+        limeSpan(kj::mv(limeSpan)) {}
+  TraceParentContext(TraceParentContext&& other) = default;
+  TraceParentContext& operator=(TraceParentContext&& other) = default;
+  KJ_DISALLOW_COPY(TraceParentContext);
+
+  SpanParent span;
+  LimeSpanParent limeSpan;
+};
+
 }  // namespace workerd
diff --git a/src/workerd/server/server.c++ b/src/workerd/server/server.c++
@@ -2138,7 +2138,8 @@ private:
         kj::Maybe<kj::String> cacheName,
         kj::Maybe<kj::String> cfBlobJson,
         SpanParent parentSpan)
-        : client(asHttpClient(parent.startRequest({kj::mv(cfBlobJson), kj::mv(parentSpan)}))),
+        : client(asHttpClient(parent.startRequest(
+              {kj::mv(cfBlobJson), TraceParentContext(kj::mv(parentSpan), nullptr)}))),
           cacheName(kj::mv(cacheName)),
           cacheNamespaceHeader(cacheNamespaceHeader) {}