@@ -109,13 +109,116 @@ class TraceId final {
109109 return high;
110110 }
111111
112+ static TraceId fromCapnp (rpc::InvocationSpanContext::TraceId::Reader reader);
113+ void toCapnp (rpc::InvocationSpanContext::TraceId::Builder writer) const ;
114+
112115private:
113116 uint64_t low = 0 ;
114117 uint64_t high = 0 ;
115118};
116119constexpr TraceId TraceId::nullId = nullptr ;
117120
121+ // The InvocationSpanContext is a tuple of a trace id, invocation id, and span id.
122+ // The trace id represents a top-level request and should be shared across all
123+ // invocation spans and events within those spans. The invocation id identifies
124+ // a specific worker invocation. The span id identifies a specific span within an
125+ // invocation. Every invocation of every worker should have an InvocationSpanContext.
126+ // That may or may not have a trigger InvocationSpanContext.
127+ class InvocationSpanContext final : public kj::Refcounted,
128+ public kj::EnableAddRefToThis<InvocationSpanContext> {
129+ public:
130+ // Spans within a InvocationSpanContext are identified by a span id that is a
131+ // monotically increasing number. Every InvocationSpanContext has a root span
132+ // whose ID is zero. Every child span context created within that context will
133+ // have a span id that is one greater than the previously created one.
134+ class SpanIdCounter final : public kj::Refcounted {
135+ public:
136+ SpanIdCounter () = default ;
137+ KJ_DISALLOW_COPY_AND_MOVE (SpanIdCounter);
138+
139+ inline kj::uint next () {
140+ #ifdef KJ_DEBUG
141+ static constexpr kj::uint kMax = kj::maxValue;
142+ KJ_ASSERT (id < kMax , " max number of spans exceeded"
143+ #endif
144+ return id++;
145+ }
146+
147+ private:
148+ kj::uint id = 1 ;
149+ };
150+
151+ // The constructor is public only so kj::rc can see it and create a new instance.
152+ // User code should use the static factory methods or the newChild method.
153+ InvocationSpanContext (kj::Badge<InvocationSpanContext>,
154+ kj::Maybe<kj::Rc<SpanIdCounter>> counter,
155+ TraceId traceId,
156+ TraceId invocationId,
157+ kj::uint spanId = 0 ,
158+ kj::Maybe<kj::Rc<InvocationSpanContext>> parentSpanContext = kj::none);
159+ KJ_DISALLOW_COPY_AND_MOVE (InvocationSpanContext);
160+
161+ inline const TraceId& getTraceId () const {
162+ return traceId;
163+ }
164+
165+ inline const TraceId& getInvocationId () const {
166+ return invocationId;
167+ }
168+
169+ inline const kj::uint getSpanId () const {
170+ return spanId;
171+ }
172+
173+ inline const kj::Maybe<kj::Rc<InvocationSpanContext>>& getParent () const {
174+ return parentSpanContext;
175+ }
176+
177+ // Creates a new child span. If the current context does not have a counter,
178+ // then this will assert. If isTrigger() is true then it will not have a
179+ // counter.
180+ kj::Rc<InvocationSpanContext> newChild ();
181+
182+ // An InvocationSpanContext is a trigger context if it has no counter. This
183+ // generally means the SpanContext was create from a capnp message and
184+ // represents an InvocationSpanContext that was propagated from a parent
185+ // or triggering context.
186+ bool isTrigger () const {
187+ return counter == kj::none;
188+ }
189+
190+ // Creates a new InvocationSpanContext. If the triggerContext is given, then its
191+ // traceId is used as the traceId for the newly created context. Otherwise a new
192+ // traceId is generated. The invocationId is always generated new and the spanId
193+ // will be 0 with no parent span.
194+ static kj::Rc<InvocationSpanContext> newForInvocation (
195+ kj::Maybe<kj::Rc<InvocationSpanContext>&> triggerContext = kj::none,
196+ kj::Maybe<kj::EntropySource&> entropySource = kj::none);
197+
198+ // Creates a new InvocationSpanContext from a capnp message. The returned
199+ // InvocationSpanContext will not be capable of creating child spans and
200+ // is considered only a "trigger" span.
201+ static kj::Maybe<kj::Rc<InvocationSpanContext>> fromCapnp (
202+ rpc::InvocationSpanContext::Reader reader);
203+ void toCapnp (rpc::InvocationSpanContext::Builder writer) const ;
204+
205+ private:
206+ // If there is no counter, then child spans cannot be created from
207+ // this InvocationSpanContext.
208+ kj::Maybe<kj::Rc<SpanIdCounter>> counter;
209+ const TraceId traceId;
210+ const TraceId invocationId;
211+ const kj::uint spanId;
212+
213+ // The parentSpanContext can be either a direct parent or a trigger
214+ // context. If it is a trigger context, then it should have the same
215+ // traceId but a different invocationId (unless predictable mode for
216+ // testing is enabled). The isTrigger() should also return true.
217+ const kj::Maybe<kj::Rc<InvocationSpanContext>> parentSpanContext;
218+ };
219+
118220kj::String KJ_STRINGIFY (const TraceId& id);
221+ kj::String KJ_STRINGIFY (const kj::Rc<InvocationSpanContext>& context);
119222} // namespace tracing
120223
121224enum class PipelineLogLevel {
0 commit comments