1
// Copyright (c) 2017-2022 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
#pragma once
6
 
7
#include <workerd/io/io-context.h>
8
#include <workerd/io/trace.h>
9
#include <workerd/util/weak-refs.h>
10
 
11
#include <kj/refcount.h>
12
 
13
namespace workerd {
14
namespace tracing {
15
class TailStreamWriter;
16
}  // namespace tracing
17
 
18
class WorkerTracer;
19
 
20
// An abstract class that defines shares functionality for tracers that have different
21
// characteristics. This interface is used to submit both buffered and streaming tail events.
22
// TODO(streaming-tail): When further consolidating the tail worker implementations, the interface
23
// of the add* methods below should make more sense: The invocation span context below is currently
24
// only being used in the streaming model, when we have switched the buffered model to streaming
25
// there will be plenty of cleanup potential.
26
class BaseTracer: public kj::Refcounted {
27
 public:
28
  virtual ~BaseTracer() noexcept(false) {
29
    selfRef->invalidate();
30
  };
31
 
32
  // Weak reference to this tracer, used by user-tracing SpanSubmitter implementations so
33
  // abandoned user promises cannot pin tracer lifetime.
34
  using WeakRef = workerd::WeakRef<BaseTracer>;
35
  kj::Own<WeakRef> getWeakRef() {
36
    return kj::addRef(*selfRef);
37
  }
38
 
39
  // Adds log line to trace.  For Spectre, timestamp should only be as accurate as JS Date.now().
40
  virtual void addLog(const tracing::InvocationSpanContext& context,
41
      kj::Date timestamp,
42
      LogLevel logLevel,
43
      kj::String message) = 0;
44
  // Add a span open event.
45
  virtual void addSpanOpen(tracing::SpanId spanId,
46
      tracing::SpanId parentSpanId,
47
      kj::ConstString operationName,
48
      kj::Date startTime) = 0;
49
  // Add a span close event.
50
  virtual void addSpanClose(tracing::SpanEndData&& span, kj::Maybe<kj::Date> maybeStartTime) = 0;
51
 
52
  virtual void addException(const tracing::InvocationSpanContext& context,
53
      kj::Date timestamp,
54
      kj::String name,
55
      kj::String message,
56
      kj::Maybe<kj::String> stack) = 0;
57
 
58
  virtual void addDiagnosticChannelEvent(const tracing::InvocationSpanContext& context,
59
      kj::Date timestamp,
60
      kj::String channel,
61
      kj::Array<kj::byte> message) = 0;
62
 
63
  // Adds info about the event that triggered the trace.  Must not be called more than once.
64
  virtual void setEventInfo(
65
      IoContext::IncomingRequest& incomingRequest, tracing::EventInfo&& info) = 0;
66
 
67
  // Sets the return event for Streaming Tail Worker, including fetchResponseInfo (HTTP status code)
68
  // if available. Must not be called more than once, and fetchResponseInfo should only be set for
69
  // fetch events. For buffered tail worker, there is no distinct return event so we only add
70
  // fetchResponseInfo to the trace if present.
71
  virtual void setReturn(kj::Maybe<kj::Date> time = kj::none,
72
      kj::Maybe<tracing::FetchResponseInfo> fetchResponseInfo = kj::none) = 0;
73
 
74
  // Reports the outcome event of the worker invocation. For Streaming Tail Worker, this will be the
75
  // final event, causing the stream to terminate.
76
  virtual void setOutcome(EventOutcome outcome, kj::Duration cpuTime, kj::Duration wallTime) = 0;
77
 
78
  // Report time as seen from the incoming Request when the request is complete, since it will not
79
  // be available afterwards.
80
  virtual void recordTimestamp(kj::Date timestamp) = 0;
81
 
82
  SpanParent makeUserRequestSpan(
83
      tracing::TraceId traceId, kj::Maybe<tracing::TraceFlags> traceFlags);
84
 
85
  using MakeUserRequestSpanFunc =
86
      kj::Function<SpanParent(tracing::TraceId, kj::Maybe<tracing::TraceFlags>)>;
87
 
88
  // Allow setting the user request span after the tracer has been created so its observer can
89
  // reference the tracer. This can only be set once.
90
  void setMakeUserRequestSpanFunc(MakeUserRequestSpanFunc func);
91
 
92
  virtual void setJsRpcInfo(const tracing::InvocationSpanContext& context,
93
      kj::Date timestamp,
94
      const kj::ConstString& methodName) = 0;
95
 
96
  // Mark this tracer as intentionally unused (e.g., for duplicate alarm requests).
97
  // When set, the destructor will not log a warning about missing Onset event.
98
  void markUnused() {
99
    markedUnused = true;
100
  }
101
 
102
 protected:
103
  // Retrieves the current timestamp. If the IoContext is no longer available, we assume that the
104
  // worker must have wrapped up and reported its outcome event, we report completeTime in that case
105
  // acordingly.
106
  kj::Date getTime();
107
 
108
  // helper method for addSpanEnd() implementations
109
  void adjustSpanTime(tracing::SpanEndData& span, kj::Maybe<kj::Date> maybeStartTime);
110
 
111
  // Function to create the root span for the new tracing format.
112
  kj::Maybe<MakeUserRequestSpanFunc> makeUserRequestSpanFunc;
113
 
114
  // Time to be reported for the outcome event time. This will be set before the outcome is
115
  // dispatched.
116
  kj::Date completeTime = kj::UNIX_EPOCH;
117
 
118
  // Weak reference to the IoContext, used to report span end time if available.
119
  kj::Maybe<kj::Own<IoContext::WeakRef>> weakIoContext;
120
 
121
  // When true, the destructor will not log a warning about missing Onset event.
122
  // Set via markUnused() when a tracer is intentionally not used (e.g., duplicate alarm requests).
123
  bool markedUnused = false;
124
 
125
 private:
126
  friend class workerd::WeakRef<BaseTracer>;
127
  kj::Own<WeakRef> selfRef = kj::refcounted<WeakRef>(kj::Badge<BaseTracer>(), *this);
128
};
129
 
130
// Records a worker stage's trace information into a Trace object.  When all references to the
131
// Tracer are released, its Trace is considered complete and ready for submission.
132
class WorkerTracer final: public BaseTracer {
133
 public:
134
  explicit WorkerTracer(kj::Maybe<kj::Rc<kj::Refcounted>> parentPipeline,
135
      kj::Own<Trace> trace,
136
      PipelineLogLevel pipelineLogLevel,
137
      kj::Maybe<kj::Array<tracing::Attribute>> tailAttributes,
138
      kj::Maybe<kj::Own<tracing::TailStreamWriter>> maybeTailStreamWriter);
139
  virtual ~WorkerTracer() noexcept(false);
140
  KJ_DISALLOW_COPY_AND_MOVE(WorkerTracer);
141
 
142
  // Returns a promise that fulfills when trace is complete. Only one such promise can
143
  // exist at a time. Used in workerd, where we don't have to worry about pipelines.
144
  kj::Promise<kj::Own<Trace>> onComplete();
145
 
146
  void addLog(const tracing::InvocationSpanContext& context,
147
      kj::Date timestamp,
148
      LogLevel logLevel,
149
      kj::String message) override;
150
  void addSpanOpen(tracing::SpanId spanId,
151
      tracing::SpanId parentSpanId,
152
      kj::ConstString operationName,
153
      kj::Date startTime) override;
154
  void addSpanClose(tracing::SpanEndData&& span, kj::Maybe<kj::Date> maybeStartTime) override;
155
  void addException(const tracing::InvocationSpanContext& context,
156
      kj::Date timestamp,
157
      kj::String name,
158
      kj::String message,
159
      kj::Maybe<kj::String> stack) override;
160
  void addDiagnosticChannelEvent(const tracing::InvocationSpanContext& context,
161
      kj::Date timestamp,
162
      kj::String channel,
163
      kj::Array<kj::byte> message) override;
164
  // Set event info (equivalent to Onset event under streaming). We use the incomingRequest here
165
  // since the IoContext may not have the IncomingRequest linked to it yet (depending on if
166
  // delivered() has been set), so it might not be possible to acquire the required timestamp and
167
  // span context from it.
168
  void setEventInfo(
169
      IoContext::IncomingRequest& incomingRequest, tracing::EventInfo&& info) override;
170
  // Variant for when we don't have a proper IoContext but instead provide context and timestamp
171
  // directly, used internally for RPC-based tracing.
172
  void setEventInfoInternal(
173
      const tracing::InvocationSpanContext& context, kj::Date timestamp, tracing::EventInfo&& info);
174
 
175
  void setOutcome(EventOutcome outcome, kj::Duration cpuTime, kj::Duration wallTime) override;
176
  virtual void recordTimestamp(kj::Date timestamp) override;
177
 
178
  // Set a worker-level tag/attribute to be provided in the onset event.
179
  void setWorkerAttribute(kj::ConstString key, Span::TagValue value);
180
 
181
  void setReturn(kj::Maybe<kj::Date> time = kj::none,
182
      kj::Maybe<tracing::FetchResponseInfo> fetchResponseInfo = kj::none) override;
183
 
184
  void setJsRpcInfo(const tracing::InvocationSpanContext& context,
185
      kj::Date timestamp,
186
      const kj::ConstString& methodName) override;
187
 
188
 private:
189
  PipelineLogLevel pipelineLogLevel;
190
  kj::Own<Trace> trace;
191
  // span attributes to be added to the onset event.
192
  kj::Vector<tracing::Attribute> attributes;
193
 
194
  // TODO(streaming-tail): Top-level invocation span context, used to add a placeholder span context
195
  // for trace events. This should no longer be needed after merging the existing span ID and
196
  // InvocationSpanContext interfaces.
197
  kj::Maybe<tracing::InvocationSpanContext> topLevelInvocationSpanContext;
198
 
199
  // own an instance of the pipeline to make sure it doesn't get destroyed
200
  // before we're finished tracing. kj::Refcounted serves as a fill-in here since the pipeline
201
  // tracer is not needed otherwise.
202
  kj::Maybe<kj::Rc<kj::Refcounted>> parentPipeline;
203
  kj::Maybe<kj::Own<kj::PromiseFulfiller<kj::Own<Trace>>>> completeFulfiller;
204
 
205
  kj::Maybe<kj::Own<tracing::TailStreamWriter>> maybeTailStreamWriter;
206
};
207
 
208
class SpanSubmitter: public kj::Refcounted {
209
 public:
210
  // Called when a span is opened. Submitters may buffer this for later assembly or stream it
211
  // immediately.
212
  virtual bool submitSpanOpen(tracing::SpanId spanId,
213
      tracing::SpanId parentSpanId,
214
      kj::ConstString operationName,
215
      kj::Date startTime) = 0;
216
 
217
  // Variant for spans opened directly by user JS (`ctx.tracing.enterSpan`). Submitters can
218
  // apply different policies (e.g. bypass operation-name allowlists for user spans).
219
  virtual bool submitUserSpanOpen(tracing::SpanId spanId,
220
      tracing::SpanId parentSpanId,
221
      kj::ConstString operationName,
222
      kj::Date startTime) {
223
    return submitSpanOpen(spanId, parentSpanId, kj::mv(operationName), startTime);
224
  }
225
 
226
  // Called when a span is closed. Together with the open data, provides all span information.
227
  virtual void submitSpanClose(
228
      tracing::SpanId spanId, kj::Date startTime, kj::Date endTime, Span::TagMap&& tags) = 0;
229
 
230
  virtual tracing::SpanId makeSpanId() = 0;
231
};
232
 
233
// The user tracing observer
234
class UserSpanObserver final: public SpanObserver {
235
 public:
236
  // constructor for top-level observer
237
  UserSpanObserver(kj::Own<SpanSubmitter> submitter)
238
      : submitter(kj::mv(submitter)),
239
        spanId(tracing::SpanId::nullId),
240
        parentSpanId(tracing::SpanId::nullId),
241
        traceId(nullptr) {}
242
  // constructor for top-level observer with trace ID and optional trace flags
243
  UserSpanObserver(kj::Own<SpanSubmitter> submitter,
244
      tracing::TraceId traceId,
245
      kj::Maybe<tracing::TraceFlags> traceFlags = kj::none)
246
      : submitter(kj::mv(submitter)),
247
        spanId(tracing::SpanId::nullId),
248
        parentSpanId(tracing::SpanId::nullId),
249
        traceId(kj::mv(traceId)),
250
        traceFlags(traceFlags) {}
251
  // constructor for subsequent observers attached to a span. `fromUserCode` is true for
252
  // spans created directly via `ctx.tracing.enterSpan`; this routes onOpen() through
253
  // submitUserSpanOpen() so submitters can apply different policies than for runtime spans.
254
  UserSpanObserver(kj::Own<SpanSubmitter> submitter,
255
      tracing::SpanId parentSpanId,
256
      tracing::TraceId traceId,
257
      kj::Maybe<tracing::TraceFlags> traceFlags,
258
      bool fromUserCode = false)
259
      : submitter(kj::mv(submitter)),
260
        spanId(this->submitter->makeSpanId()),
261
        parentSpanId(parentSpanId),
262
        traceId(kj::mv(traceId)),
263
        traceFlags(traceFlags),
264
        fromUserCode(fromUserCode) {}
265
  KJ_DISALLOW_COPY(UserSpanObserver);
266
 
267
  kj::Own<SpanObserver> newChild() override;
268
  kj::Own<SpanObserver> newChildFromUserCode() override;
269
  void onOpen(kj::ConstString operationName, kj::Date startTime) override;
270
  void onClose(kj::Date endTime, Span::TagMap&& tags, kj::Vector<Span::Log>&& logs) override;
271
  kj::Date getTime() override;
272
  kj::Maybe<tracing::SpanContext> toSpanContext() override;
273
  tracing::SpanId getSpanId() override;
274
 
275
 private:
276
  kj::Own<SpanSubmitter> submitter;
277
  tracing::SpanId spanId;
278
  tracing::SpanId parentSpanId;
279
  tracing::TraceId traceId;
280
  // Ideally, the span observer wouldn't need to maintain startTime, but we still need it at present
281
  // for timestamp adjustments.
282
  kj::Date startTime = kj::UNIX_EPOCH;
283
  // Allow the submitter to reject spans, causing them to not be reported.
284
  bool wasAccepted = true;
285
  kj::Maybe<tracing::TraceFlags> traceFlags;
286
  // True only for spans created directly via `ctx.tracing.enterSpan`. Not inherited by
287
  // children created via `newChild()` — runtime sub-operations nested inside an enterSpan
288
  // callback (e.g. `kv_get`) should remain subject to runtime-span policies.
289
  bool fromUserCode = false;
290
};
291
 
292
}  // namespace workerd