[nfc] Refactor and add trace span RPC code

This will be used to better export user tracing spans in a follow-up PR.

src/workerd/io/BUILD.bazel

@@ -156,6 +156,7 @@ wd_cc_library(
     ],
     visibility = ["//visibility:public"],
     deps = [
+        ":trace_capnp",
         ":worker-interface_capnp",
         "//src/workerd/jsg:memory-tracker",
         "//src/workerd/util",
@@ -273,6 +274,8 @@ wd_capnp_library(src = "outcome.capnp")
 
 wd_capnp_library(src = "script-version.capnp")
 
+wd_capnp_library(src = "trace.capnp")
+
 wd_capnp_library(src = "compatibility-date.capnp")
 
 wd_capnp_library(src = "features.capnp")

src/workerd/io/trace.c++

@@ -1729,28 +1729,62 @@ void WorkerTracer::addSpan(const Span& span, kj::String spanContext) {
   // TODO(cleanup): Create a function in kj::OneOf to automatically convert to a given type (i.e
   // String) to avoid having to handle each type explicitly here.
   for (const Span::TagMap::Entry& tag: span.tags) {
-    auto value = [&]() {
-      KJ_SWITCH_ONEOF(tag.value) {
-        KJ_CASE_ONEOF(str, kj::String) {
-          return kj::str(str);
-        }
-        KJ_CASE_ONEOF(val, int64_t) {
-          return kj::str(val);
-        }
-        KJ_CASE_ONEOF(val, double) {
-          return kj::str(val);
-        }
-        KJ_CASE_ONEOF(val, bool) {
-          return kj::str(val);
-        }
-      }
-      KJ_UNREACHABLE;
-    }();
-    kj::String message = kj::str("[\"tag: "_kj, tag.key, " => "_kj, value, "\"]");
+    kj::String message = kj::str("[\"tag: "_kj, tag.key, " => "_kj, spanTagStr(tag.value), "\"]");
     addLog(span.endTime, LogLevel::LOG, kj::mv(message), true);
   }
 }
 
+kj::String spanTagStr(const kj::OneOf<bool, int64_t, double, kj::String>& tag) {
+  KJ_SWITCH_ONEOF(tag) {
+    KJ_CASE_ONEOF(str, kj::String) {
+      return kj::str(str);
+    }
+    KJ_CASE_ONEOF(val, int64_t) {
+      return kj::str(val);
+    }
+    KJ_CASE_ONEOF(val, double) {
+      return kj::str(val);
+    }
+    KJ_CASE_ONEOF(val, bool) {
+      return kj::str(val);
+    }
+  }
+  KJ_UNREACHABLE;
+}
+
+using RpcValue = rpc::TagValue;
+void serializeTagValue(RpcValue::Builder builder, const Span::TagValue& value) {
+  KJ_SWITCH_ONEOF(value) {
+    KJ_CASE_ONEOF(b, bool) {
+      builder.setBool(b);
+    }
+    KJ_CASE_ONEOF(i, int64_t) {
+      builder.setInt64(i);
+    }
+    KJ_CASE_ONEOF(d, double) {
+      builder.setFloat64(d);
+    }
+    KJ_CASE_ONEOF(s, kj::String) {
+      builder.setString(s);
+    }
+  }
+}
+
+Span::TagValue deserializeTagValue(RpcValue::Reader value) {
+  switch (value.which()) {
+    case RpcValue::BOOL:
+      return value.getBool();
+    case RpcValue::FLOAT64:
+      return value.getFloat64();
+    case RpcValue::INT64:
+      return value.getInt64();
+    case RpcValue::STRING:
+      return kj::heapString(value.getString());
+    default:
+      KJ_UNREACHABLE;
+  }
+}
+
 void WorkerTracer::addException(
     kj::Date timestamp, kj::String name, kj::String message, kj::Maybe<kj::String> stack) {
   if (trace->exceededExceptionLimit) {

src/workerd/io/trace.capnp

@@ -0,0 +1,30 @@
+# Copyright (c) 2017-2022 Cloudflare, Inc.
+# Licensed under the Apache 2.0 license found in the LICENSE file or at:
+#     https://opensource.org/licenses/Apache-2.0
+
+@0xc40f73be329a38d9;
+
+using Cxx = import "/capnp/c++.capnp";
+$Cxx.namespace("workerd::rpc");
+
+# This file contains trace helper structures. The Trace struct is deliberately not defined here –
+# many files just need to have the span/tag definitions available. Since Trace contains interfaces,
+# it also causes a large amount of code to be generated in the .capnp.h file, which affects header
+# parsing overhead/compile times for every file that depends on the capnp file defining Trace.
+
+# The value of a span tag.
+struct TagValue {
+  union {
+    string @0 :Text;
+    bool @1 :Bool;
+    int64 @2 :Int64;
+    float64 @3 :Float64;
+  }
+}
+
+# A key/value span tag for tracing.
+struct Tag {
+  key @0 :Text;
+  value @1 :TagValue;
+}
+

src/workerd/io/trace.h

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <workerd/io/outcome.capnp.h>
+#include <workerd/io/trace.capnp.h>
 #include <workerd/io/worker-interface.capnp.h>
 #include <workerd/jsg/memory.h>
 #include <workerd/util/own-util.h>
@@ -94,7 +95,7 @@ class TraceId final {
   // Replicates W3C Serialization
   kj::String toW3C() const;
 
-  // Creates a random Trace Id, optionally usig a given entropy source. If an
+  // Creates a random Trace Id, optionally using a given entropy source. If an
   // entropy source is not given, then we fallback to using BoringSSL's RAND_bytes.
   static TraceId fromEntropy(kj::Maybe<kj::EntropySource&> entropy = kj::none);
 
@@ -977,7 +978,7 @@ class WorkerTracer final: public kj::Refcounted, public BaseTracer {
   // before we're finished tracing
   kj::Maybe<kj::Rc<PipelineTracer>> parentPipeline;
   // A weak reference for the internal span submitter. We use this so that the span submitter can
-  // add spans while the tracer exists, but does not artifically prolong the lifetime of the tracer
+  // add spans while the tracer exists, but does not artificially prolong the lifetime of the tracer
   // which would interfere with span submission (traces get submitted when the worker returns its
   // response, but with e.g. waitUntil() the worker can still be performing tasks afterwards so the
   // span submitter may exist for longer than the tracer).
@@ -1044,6 +1045,12 @@ struct Span {
         endTime(startTime) {}
 };
 
+// Utility functions for handling span tags.
+void serializeTagValue(rpc::TagValue::Builder builder, const Span::TagValue& value);
+Span::TagValue deserializeTagValue(rpc::TagValue::Reader value);
+// Stringifier for span tags, getting this to work with KJ_STRINGIFY() appears exceedingly difficult.
+kj::String spanTagStr(const kj::OneOf<bool, int64_t, double, kj::String>& tag);
+
 // An opaque token which can be used to create child spans of some parent. This is typically
 // passed down from a caller to a callee when the caller wants to allow the callee to create
 // spans for itself that show up as children of the caller's span, but the caller does not