Skip to content

Commit

Permalink
InvocationSpanContext doesn't need to be refcounted
Browse files Browse the repository at this point in the history
Now that InvocationSpanContext is not using a stateful counter
internally there's no reason for it to be refcounted. Simplify!
  • Loading branch information
jasnell committed Dec 6, 2024
1 parent f845ad5 commit 4d962ab
Show file tree
Hide file tree
Showing 7 changed files with 145 additions and 92 deletions.
2 changes: 1 addition & 1 deletion src/workerd/io/io-context.c++
Original file line number Diff line number Diff line change
Expand Up @@ -202,7 +202,7 @@ IoContext::IncomingRequest::IoContext_IncomingRequest(kj::Own<IoContext> context
kj::Own<IoChannelFactory> ioChannelFactoryParam,
kj::Own<RequestObserver> metricsParam,
kj::Maybe<kj::Own<WorkerTracer>> workerTracer,
kj::Rc<tracing::InvocationSpanContext> invocationSpanContext)
tracing::InvocationSpanContext invocationSpanContext)
: context(kj::mv(contextParam)),
metrics(kj::mv(metricsParam)),
workerTracer(kj::mv(workerTracer)),
Expand Down
6 changes: 3 additions & 3 deletions src/workerd/io/io-context.h
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ class IoContext_IncomingRequest final {
kj::Own<IoChannelFactory> ioChannelFactory,
kj::Own<RequestObserver> metrics,
kj::Maybe<kj::Own<WorkerTracer>> workerTracer,
kj::Rc<tracing::InvocationSpanContext> invocationSpanContext);
tracing::InvocationSpanContext invocationSpanContext);
KJ_DISALLOW_COPY_AND_MOVE(IoContext_IncomingRequest);
~IoContext_IncomingRequest() noexcept(false);

Expand Down Expand Up @@ -163,7 +163,7 @@ class IoContext_IncomingRequest final {

// The invocation span context is a unique identifier for a specific
// worker invocation.
kj::Rc<tracing::InvocationSpanContext>& getInvocationSpanContext() {
tracing::InvocationSpanContext& getInvocationSpanContext() {
return invocationSpanContext;
}

Expand All @@ -176,7 +176,7 @@ class IoContext_IncomingRequest final {
// The invocation span context identifies the trace id, invocation id, and root
// span for the current request. Every invocation of a worker function always
// has a root span, even if it is not explicitly traced.
kj::Rc<tracing::InvocationSpanContext> invocationSpanContext;
tracing::InvocationSpanContext invocationSpanContext;

bool wasDelivered = false;

Expand Down
56 changes: 28 additions & 28 deletions src/workerd/io/trace-test.c++
Original file line number Diff line number Diff line change
Expand Up @@ -91,46 +91,46 @@ KJ_TEST("InvocationSpanContext") {

// We can create an InvocationSpanContext...
static constexpr auto kCheck = TraceId(0x2a2a2a2a2a2a2a2a, 0x2a2a2a2a2a2a2a2a);
KJ_EXPECT(sc->getTraceId() == kCheck);
KJ_EXPECT(sc->getInvocationId() == kCheck);
KJ_EXPECT(sc->getSpanId() == SpanId(1));
KJ_EXPECT(sc.getTraceId() == kCheck);
KJ_EXPECT(sc.getInvocationId() == kCheck);
KJ_EXPECT(sc.getSpanId() == SpanId(1));

// And serialize that to a capnp struct...
capnp::MallocMessageBuilder builder;
auto root = builder.initRoot<rpc::InvocationSpanContext>();
sc->toCapnp(root);
sc.toCapnp(root);

// Then back again...
auto sc2 = KJ_ASSERT_NONNULL(InvocationSpanContext::fromCapnp(root.asReader()));
KJ_EXPECT(sc2->getTraceId() == kCheck);
KJ_EXPECT(sc2->getInvocationId() == kCheck);
KJ_EXPECT(sc2->getSpanId() == SpanId(1));
KJ_EXPECT(sc2->isTrigger());
KJ_EXPECT(sc2.getTraceId() == kCheck);
KJ_EXPECT(sc2.getInvocationId() == kCheck);
KJ_EXPECT(sc2.getSpanId() == SpanId(1));
KJ_EXPECT(sc2.isTrigger());

// The one that has been deserialized from capnp cannot create children...
try {
sc2->newChild();
sc2.newChild();
KJ_FAIL_ASSERT("should not be able to create child span with SpanContext from capnp");
} catch (kj::Exception& ex) {
KJ_EXPECT(ex.getDescription() ==
"expected !isTrigger(); unable to create child spans on this context"_kj);
}

auto sc3 = sc->newChild();
KJ_EXPECT(sc3->getTraceId() == kCheck);
KJ_EXPECT(sc3->getInvocationId() == kCheck);
KJ_EXPECT(sc3->getSpanId() == SpanId(2));
auto sc3 = sc.newChild();
KJ_EXPECT(sc3.getTraceId() == kCheck);
KJ_EXPECT(sc3.getInvocationId() == kCheck);
KJ_EXPECT(sc3.getSpanId() == SpanId(2));

auto sc4 = InvocationSpanContext::newForInvocation(sc2, fakeEntropySource);
KJ_EXPECT(sc4->getTraceId() == kCheck);
KJ_EXPECT(sc4->getInvocationId() == kCheck);
KJ_EXPECT(sc4->getSpanId() == SpanId(3));

auto& sc5 = KJ_ASSERT_NONNULL(sc4->getParent());
KJ_EXPECT(sc5->getTraceId() == kCheck);
KJ_EXPECT(sc5->getInvocationId() == kCheck);
KJ_EXPECT(sc5->getSpanId() == SpanId(1));
KJ_EXPECT(sc5->isTrigger());
KJ_EXPECT(sc4.getTraceId() == kCheck);
KJ_EXPECT(sc4.getInvocationId() == kCheck);
KJ_EXPECT(sc4.getSpanId() == SpanId(3));

auto& sc5 = KJ_ASSERT_NONNULL(sc4.getParent());
KJ_EXPECT(sc5.getTraceId() == kCheck);
KJ_EXPECT(sc5.getInvocationId() == kCheck);
KJ_EXPECT(sc5.getSpanId() == SpanId(1));
KJ_EXPECT(sc5.isTrigger());
}

KJ_TEST("Read/Write FetchEventInfo works") {
Expand Down Expand Up @@ -532,9 +532,9 @@ KJ_TEST("Read/Write TraceEvent works") {
tracing::TailEvent info2(reader);
KJ_ASSERT(info2.timestamp == kj::UNIX_EPOCH);
KJ_ASSERT(info2.sequence == 0);
KJ_ASSERT(info2.context.invocationId == context->getInvocationId());
KJ_ASSERT(info2.context.traceId == context->getTraceId());
KJ_ASSERT(info2.context.spanId == context->getSpanId());
KJ_ASSERT(info2.context.invocationId == context.getInvocationId());
KJ_ASSERT(info2.context.traceId == context.getTraceId());
KJ_ASSERT(info2.context.spanId == context.getSpanId());

auto& event = KJ_ASSERT_NONNULL(info2.event.tryGet<tracing::Mark>());
auto& log2 = KJ_ASSERT_NONNULL(event.tryGet<tracing::Log>());
Expand All @@ -545,9 +545,9 @@ KJ_TEST("Read/Write TraceEvent works") {
tracing::TailEvent info3 = info.clone();
KJ_ASSERT(info3.timestamp == kj::UNIX_EPOCH);
KJ_ASSERT(info3.sequence == 0);
KJ_ASSERT(info3.context.invocationId == context->getInvocationId());
KJ_ASSERT(info3.context.traceId == context->getTraceId());
KJ_ASSERT(info3.context.spanId == context->getSpanId());
KJ_ASSERT(info3.context.invocationId == context.getInvocationId());
KJ_ASSERT(info3.context.traceId == context.getTraceId());
KJ_ASSERT(info3.context.spanId == context.getSpanId());

auto& event2 = KJ_ASSERT_NONNULL(info3.event.tryGet<tracing::Mark>());
auto& log3 = KJ_ASSERT_NONNULL(event2.tryGet<tracing::Log>());
Expand Down
85 changes: 49 additions & 36 deletions src/workerd/io/trace.c++
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ kj::String TraceId::toW3C() const {
}

namespace {
uint64_t getRandom64Bit(kj::Maybe<kj::EntropySource&>& entropySource) {
uint64_t getRandom64Bit(const kj::Maybe<kj::EntropySource&>& entropySource) {
uint64_t ret = 0;
uint8_t tries = 0;

Expand Down Expand Up @@ -178,31 +178,34 @@ InvocationSpanContext::InvocationSpanContext(kj::Badge<InvocationSpanContext>,
TraceId traceId,
TraceId invocationId,
SpanId spanId,
kj::Maybe<kj::Rc<InvocationSpanContext>> parentSpanContext)
kj::Maybe<const InvocationSpanContext&> parentSpanContext)
: entropySource(entropySource),
traceId(kj::mv(traceId)),
invocationId(kj::mv(invocationId)),
spanId(kj::mv(spanId)),
parentSpanContext(kj::mv(parentSpanContext)) {}
parentSpanContext(parentSpanContext.map([](const InvocationSpanContext& ctx) {
return kj::heap<InvocationSpanContext>(ctx.clone());
})) {}

kj::Rc<InvocationSpanContext> InvocationSpanContext::newChild() {
InvocationSpanContext InvocationSpanContext::newChild() const {
KJ_ASSERT(!isTrigger(), "unable to create child spans on this context");
return kj::rc<InvocationSpanContext>(kj::Badge<InvocationSpanContext>(), entropySource, traceId,
invocationId, SpanId::fromEntropy(entropySource), addRefToThis());
kj::Maybe<kj::EntropySource&> otherEntropySource = entropySource.map(
[](auto& es) -> kj::EntropySource& { return const_cast<kj::EntropySource&>(es); });
return InvocationSpanContext(kj::Badge<InvocationSpanContext>(), otherEntropySource, traceId,
invocationId, SpanId::fromEntropy(otherEntropySource), *this);
}

kj::Rc<InvocationSpanContext> InvocationSpanContext::newForInvocation(
kj::Maybe<kj::Rc<InvocationSpanContext>&> triggerContext,
InvocationSpanContext InvocationSpanContext::newForInvocation(
kj::Maybe<const InvocationSpanContext&> triggerContext,
kj::Maybe<kj::EntropySource&> entropySource) {
kj::Maybe<kj::Rc<InvocationSpanContext>> parent;
kj::Maybe<const InvocationSpanContext&> parent;
auto traceId = triggerContext
.map([&](kj::Rc<InvocationSpanContext>& ctx) {
parent = ctx.addRef();
return ctx->traceId;
.map([&](auto& ctx) mutable -> TraceId {
parent = ctx;
return ctx.traceId;
}).orDefault([&] { return TraceId::fromEntropy(entropySource); });
return kj::rc<InvocationSpanContext>(kj::Badge<InvocationSpanContext>(), entropySource,
kj::mv(traceId), TraceId::fromEntropy(entropySource), SpanId::fromEntropy(entropySource),
kj::mv(parent));
return InvocationSpanContext(kj::Badge<InvocationSpanContext>(), entropySource, kj::mv(traceId),
TraceId::fromEntropy(entropySource), SpanId::fromEntropy(entropySource), kj::mv(parent));
}

TraceId TraceId::fromCapnp(rpc::InvocationSpanContext::TraceId::Reader reader) {
Expand All @@ -214,31 +217,38 @@ void TraceId::toCapnp(rpc::InvocationSpanContext::TraceId::Builder writer) const
writer.setHigh(high);
}

kj::Maybe<kj::Rc<InvocationSpanContext>> InvocationSpanContext::fromCapnp(
kj::Maybe<InvocationSpanContext> InvocationSpanContext::fromCapnp(
rpc::InvocationSpanContext::Reader reader) {
if (!reader.hasTraceId() || !reader.hasInvocationId()) {
// If the reader does not have a traceId or invocationId field then it is
// invalid and we will just ignore it.
return kj::none;
}

auto sc = kj::rc<InvocationSpanContext>(kj::Badge<InvocationSpanContext>(), kj::none,
auto sc = InvocationSpanContext(kj::Badge<InvocationSpanContext>(), kj::none,
TraceId::fromCapnp(reader.getTraceId()), TraceId::fromCapnp(reader.getInvocationId()),
reader.getSpanId());
// If the traceId or invocationId are invalid, then we'll ignore them.
if (!sc->getTraceId() || !sc->getInvocationId()) return kj::none;
if (!sc.getTraceId() || !sc.getInvocationId()) return kj::none;
return kj::mv(sc);
}

void InvocationSpanContext::toCapnp(rpc::InvocationSpanContext::Builder writer) const {
traceId.toCapnp(writer.initTraceId());
invocationId.toCapnp(writer.initInvocationId());
writer.setSpanId(spanId);
kj::mv(getParent()); // Just invalidating the parent. Not moving it anywhere.
}

kj::String KJ_STRINGIFY(const kj::Rc<InvocationSpanContext>& context) {
return kj::str(context->getTraceId(), "-", context->getInvocationId(), "-", context->getSpanId());
InvocationSpanContext InvocationSpanContext::clone() const {
kj::Maybe<kj::EntropySource&> otherEntropySource = entropySource.map(
[](auto& es) -> kj::EntropySource& { return const_cast<kj::EntropySource&>(es); });
return InvocationSpanContext(kj::Badge<InvocationSpanContext>(), otherEntropySource, traceId,
invocationId, spanId,
parentSpanContext.map([](auto& ctx) -> const InvocationSpanContext& { return *ctx.get(); }));
}

kj::String KJ_STRINGIFY(const InvocationSpanContext& context) {
return kj::str(context.getTraceId(), "-", context.getInvocationId(), "-", context.getSpanId());
}

} // namespace tracing
Expand Down Expand Up @@ -851,7 +861,7 @@ tracing::Attribute::Attribute(kj::String name, kj::Array<Value>&& value)

namespace {
kj::OneOf<tracing::Attribute::Value, kj::Array<tracing::Attribute::Value>> readValues(
rpc::Trace::Attribute::Reader& reader) {
const rpc::Trace::Attribute::Reader& reader) {
static auto readValue =
[](rpc::Trace::Attribute::Value::Reader reader) -> tracing::Attribute::Value {
auto inner = reader.getInner();
Expand Down Expand Up @@ -955,7 +965,7 @@ tracing::Attribute tracing::Attribute::clone() {
tracing::Return::Return(kj::Maybe<tracing::Return::Info> info): info(kj::mv(info)) {}

namespace {
kj::Maybe<tracing::Return::Info> readReturnInfo(rpc::Trace::Return::Reader& reader) {
kj::Maybe<tracing::Return::Info> readReturnInfo(const rpc::Trace::Return::Reader& reader) {
auto info = reader.getInfo();
switch (info.which()) {
case rpc::Trace::Return::Info::EMPTY:
Expand Down Expand Up @@ -1015,7 +1025,7 @@ tracing::SpanOpen::SpanOpen(kj::Maybe<kj::String> operationName, kj::Maybe<Info>
info(kj::mv(info)) {}

namespace {
kj::Maybe<kj::String> readSpanOpenOperationName(rpc::Trace::SpanOpen::Reader& reader) {
kj::Maybe<kj::String> readSpanOpenOperationName(const rpc::Trace::SpanOpen::Reader& reader) {
if (!reader.hasOperationName()) return kj::none;
return kj::str(reader.getOperationName());
}
Expand Down Expand Up @@ -1144,25 +1154,29 @@ tracing::Onset::Info getInfoFromReader(const rpc::Trace::Onset::Reader& reader)
}
KJ_UNREACHABLE;
}

kj::Maybe<kj::String> getScriptNameFromReader(const rpc::Trace::Onset::Reader& reader) {
if (reader.hasScriptName()) {
return kj::str(reader.getScriptName());
}
return kj::none;
}

kj::Maybe<kj::Own<ScriptVersion::Reader>> getScriptVersionFromReader(
const rpc::Trace::Onset::Reader& reader) {
if (reader.hasScriptVersion()) {
return capnp::clone(reader.getScriptVersion());
}
return kj::none;
}

kj::Maybe<kj::String> getDispatchNamespaceFromReader(const rpc::Trace::Onset::Reader& reader) {
if (reader.hasDispatchNamespace()) {
return kj::str(reader.getDispatchNamespace());
}
return kj::none;
}

kj::Maybe<kj::Array<kj::String>> getScriptTagsFromReader(const rpc::Trace::Onset::Reader& reader) {
if (reader.hasScriptTags()) {
auto tags = reader.getScriptTags();
Expand All @@ -1174,6 +1188,7 @@ kj::Maybe<kj::Array<kj::String>> getScriptTagsFromReader(const rpc::Trace::Onset
}
return kj::none;
}

kj::Maybe<kj::String> getEntrypointFromReader(const rpc::Trace::Onset::Reader& reader) {
if (reader.hasEntryPoint()) {
return kj::str(reader.getEntryPoint());
Expand Down Expand Up @@ -1326,17 +1341,15 @@ tracing::Outcome tracing::Outcome::clone() {
}

namespace {
tracing::TailEvent::Context getContextFromSpan(
const kj::Rc<tracing::InvocationSpanContext>& context) {
tracing::TailEvent::Context getContextFromSpan(const tracing::InvocationSpanContext& context) {
return tracing::TailEvent::Context(
context->getTraceId(), context->getInvocationId(), context->getSpanId());
context.getTraceId(), context.getInvocationId(), context.getSpanId());
}

kj::Maybe<tracing::TailEvent::Context> getParentContextFromSpan(
kj::Rc<tracing::InvocationSpanContext>& context) {
return context->getParent().map([](const kj::Rc<tracing::InvocationSpanContext>& context) {
return getContextFromSpan(context);
});
const tracing::InvocationSpanContext& context) {
return context.getParent().map(
[](const tracing::InvocationSpanContext& context) { return getContextFromSpan(context); });
}
} // namespace

Expand All @@ -1360,7 +1373,7 @@ tracing::TailEvent::Context tracing::TailEvent::Context::clone() {
return Context(traceId, invocationId, spanId);
}

tracing::TailEvent::TailEvent(kj::Rc<tracing::InvocationSpanContext>& context,
tracing::TailEvent::TailEvent(const tracing::InvocationSpanContext& context,
kj::Date timestamp,
kj::uint sequence,
Event&& event)
Expand All @@ -1382,18 +1395,18 @@ tracing::TailEvent::TailEvent(Context context,
event(kj::mv(event)) {}

namespace {
tracing::TailEvent::Context readContextFromTailEvent(rpc::Trace::TailEvent::Reader& reader) {
tracing::TailEvent::Context readContextFromTailEvent(const rpc::Trace::TailEvent::Reader& reader) {
return tracing::TailEvent::Context(reader.getContext());
}

kj::Maybe<tracing::TailEvent::Context> readParentContextFromTailEvent(
rpc::Trace::TailEvent::Reader& reader) {
const rpc::Trace::TailEvent::Reader& reader) {
if (!reader.hasParentContext()) return kj::none;
return tracing::TailEvent::Context(reader.getParentContext());
}

tracing::TailEvent::Event readEventFromTailEvent(rpc::Trace::TailEvent::Reader& reader) {
auto event = reader.getEvent();
tracing::TailEvent::Event readEventFromTailEvent(const rpc::Trace::TailEvent::Reader& reader) {
const auto event = reader.getEvent();
switch (event.which()) {
case rpc::Trace::TailEvent::Event::ONSET:
return tracing::Onset(event.getOnset());
Expand Down
Loading

0 comments on commit 4d962ab

Please sign in to comment.