Use a builder for tracing::Attribute to improve ergonomics

cloudflare · Dec 6, 2024 · ac06ee8 · ac06ee8
1 parent 3590bec
commit ac06ee8
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 13 deletions.
diff --git a/src/workerd/io/trace-test.c++ b/src/workerd/io/trace-test.c++
@@ -409,14 +409,14 @@ KJ_TEST("Read/Write Attribute works") {
   capnp::MallocMessageBuilder builder;
   auto infoBuilder = builder.initRoot<rpc::Trace::Attribute>();
 
-  tracing::Attribute info(kj::str("foo"), tracing::Attribute::Value((double)123.0));
-  info.copyTo(infoBuilder);
+  tracing::Attribute::Builder("foo", 1).add(123.0).add(true).finish().copyTo(infoBuilder);
 
   auto reader = infoBuilder.asReader();
   tracing::Attribute info2(reader);
   KJ_ASSERT(info2.name == "foo"_kj);
-  auto& val = KJ_ASSERT_NONNULL(info2.value.tryGet<tracing::Attribute::Value>());
-  KJ_ASSERT(KJ_ASSERT_NONNULL(val.tryGet<double>()) == 123.0);
+  auto& vals = KJ_ASSERT_NONNULL(info2.value.tryGet<tracing::Attribute::Values>());
+  KJ_ASSERT(KJ_ASSERT_NONNULL(vals[0].tryGet<double>()) == 123.0);
+  KJ_ASSERT(KJ_ASSERT_NONNULL(vals[1].tryGet<bool>()));
 }
 
 KJ_TEST("Read/Write Return works") {

diff --git a/src/workerd/io/trace.c++ b/src/workerd/io/trace.c++
@@ -841,7 +841,11 @@ tracing::Hibernate tracing::Hibernate::clone() {
   return Hibernate();
 }
 
-tracing::Attribute::Attribute(kj::String name, kj::OneOf<Value, kj::Array<Value>>&& value)
+tracing::Attribute::Attribute(kj::String name, Value&& value)
+    : name(kj::mv(name)),
+      value(kj::mv(value)) {}
+
+tracing::Attribute::Attribute(kj::String name, kj::Array<Value>&& value)
     : name(kj::mv(name)),
       value(kj::mv(value)) {}
 
@@ -891,25 +895,30 @@ void tracing::Attribute::copyTo(rpc::Trace::Attribute::Builder builder) {
   static auto writeValue = [](auto builder, const auto& value) mutable {
     KJ_SWITCH_ONEOF(value) {
       KJ_CASE_ONEOF(str, kj::String) {
-        builder[0].initInner().setText(str.asPtr());
+        builder.initInner().setText(str.asPtr());
       }
       KJ_CASE_ONEOF(b, bool) {
-        builder[0].initInner().setBool(b);
+        builder.initInner().setBool(b);
       }
       KJ_CASE_ONEOF(f, double) {
-        builder[0].initInner().setFloat(f);
+        builder.initInner().setFloat(f);
       }
       KJ_CASE_ONEOF(i, uint32_t) {
-        builder[0].initInner().setInt(i);
+        builder.initInner().setInt(i);
       }
     }
   };
   builder.setName(name.asPtr());
   KJ_SWITCH_ONEOF(value) {
     KJ_CASE_ONEOF(value, Value) {
-      writeValue(builder.initValue(1), value);
+      writeValue(builder.initValue(1)[0], value);
+    }
+    KJ_CASE_ONEOF(values, kj::Array<Value>) {
+      auto vec = builder.initValue(values.size());
+      for (size_t n = 0; n < values.size(); n++) {
+        writeValue(vec[n], values[n]);
+      }
     }
-    KJ_CASE_ONEOF(values, kj::Array<Value>) {}
   }
 }
 

diff --git a/src/workerd/io/trace.h b/src/workerd/io/trace.h
@@ -526,19 +526,69 @@ using EventInfo = kj::OneOf<FetchEventInfo,
 // Modeled after https://opentelemetry.io/docs/concepts/signals/traces/#attributes
 struct Attribute final {
   using Value = kj::OneOf<kj::String, bool, double, uint32_t>;
+  using Values = kj::Array<Value>;
 
-  explicit Attribute(kj::String name, kj::OneOf<Value, kj::Array<Value>>&& value);
+  explicit Attribute(kj::String name, Value&& value);
+  explicit Attribute(kj::String name, Values&& values);
   Attribute(rpc::Trace::Attribute::Reader reader);
   Attribute(Attribute&&) = default;
   Attribute& operator=(Attribute&&) = default;
   KJ_DISALLOW_COPY(Attribute);
 
   kj::String name;
 
-  kj::OneOf<Value, kj::Array<Value>> value;
+  kj::OneOf<Value, Values> value;
 
   void copyTo(rpc::Trace::Attribute::Builder builder);
   Attribute clone();
+
+  class Builder final {
+   public:
+    Builder(kj::String name): name(kj::mv(name)) {}
+    Builder(kj::String name, size_t n): name(kj::mv(name)), vec(n) {}
+    Builder(kj::StringPtr name): Builder(kj::str(name)) {}
+    Builder(kj::StringPtr name, size_t n): name(kj::str(name)), vec(n) {}
+    KJ_DISALLOW_COPY_AND_MOVE(Builder);
+
+    Builder& add(kj::String value) KJ_LIFETIMEBOUND {
+      vec.add(kj::mv(value));
+      return *this;
+    }
+
+    Builder& add(kj::StringPtr value) KJ_LIFETIMEBOUND {
+      vec.add(kj::str(value));
+      return *this;
+    }
+
+    Builder& add(bool value) KJ_LIFETIMEBOUND {
+      vec.add(value);
+      return *this;
+    }
+
+    Builder& add(double value) KJ_LIFETIMEBOUND {
+      vec.add(value);
+      return *this;
+    }
+
+    Builder& add(uint32_t value) KJ_LIFETIMEBOUND {
+      vec.add(value);
+      return *this;
+    }
+
+    Attribute finish() {
+      KJ_ASSERT(vec.size() >= 1);
+      if (vec.size() == 1) {
+        auto val = kj::mv(vec[0]);
+        vec.clear();
+        return Attribute(kj::mv(name), kj::mv(val));
+      }
+      return Attribute(kj::mv(name), vec.releaseAsArray());
+    }
+
+   private:
+    kj::String name;
+    kj::Vector<Value> vec;
+  };
 };
 using CustomInfo = kj::Array<Attribute>;