Support weak references to WorkerTracer

cloudflare · Sep 13, 2024 · e34cfb5 · e34cfb5
1 parent 081f031
commit e34cfb5
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 2 deletions.
diff --git a/src/workerd/io/BUILD.bazel b/src/workerd/io/BUILD.bazel
@@ -140,6 +140,7 @@ wd_cc_library(
     deps = [
         ":worker-interface_capnp",
         "//src/workerd/jsg:memory-tracker",
+        "//src/workerd/util",
         "//src/workerd/util:own-util",
         "@capnp-cpp//src/capnp:capnp-rpc",
         "@capnp-cpp//src/capnp:capnpc",

diff --git a/src/workerd/io/trace.c++ b/src/workerd/io/trace.c++
@@ -603,11 +603,13 @@ WorkerTracer::WorkerTracer(
     kj::Own<PipelineTracer> parentPipeline, kj::Own<Trace> trace, PipelineLogLevel pipelineLogLevel)
     : pipelineLogLevel(pipelineLogLevel),
       trace(kj::mv(trace)),
-      parentPipeline(kj::mv(parentPipeline)) {}
+      parentPipeline(kj::mv(parentPipeline)),
+      weakRef(kj::refcounted<WeakRef<WorkerTracer>>(kj::Badge<WorkerTracer>{}, *this)) {}
 WorkerTracer::WorkerTracer(PipelineLogLevel pipelineLogLevel)
     : pipelineLogLevel(pipelineLogLevel),
       trace(kj::refcounted<Trace>(
-          kj::none, kj::none, kj::none, kj::none, kj::none, nullptr, kj::none)) {}
+          kj::none, kj::none, kj::none, kj::none, kj::none, nullptr, kj::none)),
+      weakRef(kj::refcounted<WeakRef<WorkerTracer>>(kj::Badge<WorkerTracer>{}, *this)) {}
 
 void WorkerTracer::log(kj::Date timestamp, LogLevel logLevel, kj::String message, bool isSpan) {
   if (trace->exceededLogLimit) {

diff --git a/src/workerd/io/trace.h b/src/workerd/io/trace.h
@@ -14,6 +14,7 @@
 #include <workerd/io/outcome.capnp.h>
 #include <workerd/io/worker-interface.capnp.h>
 #include <workerd/util/own-util.h>
+#include <workerd/util/weak-refs.h>
 #include <workerd/jsg/memory.h>
 
 namespace kj {
@@ -377,6 +378,9 @@ class WorkerTracer final: public kj::Refcounted {
       kj::Own<Trace> trace,
       PipelineLogLevel pipelineLogLevel);
   explicit WorkerTracer(PipelineLogLevel pipelineLogLevel);
+  ~WorkerTracer() {
+    weakRef->invalidate();
+  }
   KJ_DISALLOW_COPY_AND_MOVE(WorkerTracer);
 
   // Adds log line to trace.  For Spectre, timestamp should only be as accurate as JS Date.now().
@@ -413,13 +417,18 @@ class WorkerTracer final: public kj::Refcounted {
   // parent process after receiving a trace from a process sandbox.
   void setTrace(rpc::Trace::Reader reader);
 
+  kj::Own<WeakRef<WorkerTracer>> addWeakRef() {
+    return weakRef->addRef();
+  }
+
 private:
   PipelineLogLevel pipelineLogLevel;
   kj::Own<Trace> trace;
 
   // own an instance of the pipeline to make sure it doesn't get destroyed
   // before we're finished tracing
   kj::Maybe<kj::Own<PipelineTracer>> parentPipeline;
+  kj::Own<WeakRef<WorkerTracer>> weakRef;
 };
 
 // =======================================================================================