Merge pull request #3740 from cloudflare/maizatskyi/2025-03-17-atomicref

kj::Arc<Worker>

Author: mikea

src/workerd/api/cache.c++

@@ -553,7 +553,7 @@ kj::Own<kj::HttpClient> Cache::getHttpClient(IoContext& context,
     .featureFlagsForFl = kj::none,
   };
   if (enableCompatFlags) {
-    metadata.featureFlagsForFl = context.getWorker().getIsolate().getFeatureFlagsForFl();
+    metadata.featureFlagsForFl = context.getWorker()->getIsolate().getFeatureFlagsForFl();
   }
   auto httpClient =
       cacheName.map([&](kj::String& n) {

src/workerd/api/queue.c++

@@ -600,7 +600,7 @@ kj::Promise<WorkerInterface::CustomEvent::Result> QueueCustomEventImpl::run(
   //   returns and the promise that it returns (if any) has resolved.
   // * In the disabled path, the queue event isn't complete until all waitUntil'ed promises resolve.
   //   This was how Queues originally worked, but made for a poor user experience.
-  auto compatFlags = context.getWorker().getIsolate().getApi().getFeatureFlags();
+  auto compatFlags = context.getWorker()->getIsolate().getApi().getFeatureFlags();
   if (compatFlags.getQueueConsumerNoWaitForWaitUntil()) {
     // The user has opted in to only waiting on their event handler rather than all waitUntil'd
     // promises.
@@ -691,7 +691,7 @@ kj::Promise<WorkerInterface::CustomEvent::Result> QueueCustomEventImpl::run(
         }
       }
       auto& ioContext = incomingRequest->getContext();
-      auto scriptId = ioContext.getWorker().getScript().getId();
+      auto scriptId = ioContext.getWorker()->getScript().getId();
       auto tasks = ioContext.getWaitUntilTasks().trace();
       if (result == IoContext_IncomingRequest::FinishScheduledResult::TIMEOUT) {
         KJ_LOG(WARNING, "NOSENTRY queue event hit timeout", scriptId, status, tasks);

src/workerd/io/io-context.c++

@@ -125,7 +125,7 @@ class IoContext::TimeoutManagerImpl::TimeoutState {
 };
 
 IoContext::IoContext(ThreadContext& thread,
-    kj::Own<const Worker> workerParam,
+    kj::Arc<Worker> workerParam,
     kj::Maybe<Worker::Actor&> actorParam,
     kj::Own<LimitEnforcer> limitEnforcerParam)
     : thread(thread),
@@ -1382,7 +1382,7 @@ kj::Promise<void> IoContext::startDeleteQueueSignalTask(IoContext* context) {
 // ======================================================================================
 
 WarningAggregator::WarningAggregator(IoContext& context, EmitCallback emitter)
-    : worker(kj::atomicAddRef(context.getWorker())),
+    : worker(context.getWorker().addRef()),
       requestMetrics(kj::addRef(context.getMetrics())),
       emitter(kj::mv(emitter)) {}
 

src/workerd/io/io-context.h

@@ -87,7 +87,7 @@ class WarningAggregator final: public kj::AtomicRefcounted {
   using Map = kj::HashMap<const Key&, kj::Own<WarningAggregator>>;
 
  private:
-  kj::Own<const Worker> worker;
+  kj::Arc<Worker> worker;
   kj::Own<RequestObserver> requestMetrics;
   EmitCallback emitter;
   kj::MutexGuarded<kj::Vector<kj::Own<WarningContext>>> warnings;
@@ -224,7 +224,7 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler
 
   // Construct a new IoContext. Before using it, you must also create an IncomingRequest.
   IoContext(ThreadContext& thread,
-      kj::Own<const Worker> worker,
+      kj::Arc<Worker> worker,
       kj::Maybe<Worker::Actor&> actor,
       kj::Own<LimitEnforcer> limitEnforcer);
 
@@ -233,8 +233,8 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler
 
   using IncomingRequest = IoContext_IncomingRequest;
 
-  const Worker& getWorker() {
-    return *worker;
+  const kj::Arc<Worker>& getWorker() {
+    return worker;
   }
   Worker::Lock& getCurrentLock() {
     return KJ_REQUIRE_NONNULL(currentLock);
@@ -835,7 +835,7 @@ class IoContext final: public kj::Refcounted, private kj::TaskSet::ErrorHandler
 
   kj::Own<WeakRef> selfRef = kj::refcounted<WeakRef>(kj::Badge<IoContext>(), *this);
 
-  kj::Own<const Worker> worker;
+  kj::Arc<Worker> worker;
   kj::Maybe<Worker::Actor&> actor;
   kj::Own<LimitEnforcer> limitEnforcer;
 

src/workerd/io/worker-entrypoint.c++

@@ -43,7 +43,7 @@ class WorkerEntrypoint final: public WorkerInterface {
   // topLevelRequest.getZoneDefaultWorkerLimits(), since the top level request may be shared between
   // zone and non-zone workers.
   static kj::Own<WorkerInterface> construct(ThreadContext& threadContext,
-      kj::Own<const Worker> worker,
+      kj::Arc<Worker> worker,
       kj::Maybe<kj::StringPtr> entrypointName,
       Frankenvalue props,
       kj::Maybe<kj::Own<Worker::Actor>> actor,
@@ -93,7 +93,7 @@ class WorkerEntrypoint final: public WorkerInterface {
   kj::Maybe<kj::Own<WorkerInterface>> failOpenService;
   bool loggedExceptionEarlier = false;
 
-  void init(kj::Own<const Worker> worker,
+  void init(kj::Arc<Worker> worker,
       kj::Maybe<kj::Own<Worker::Actor>> actor,
       kj::Own<LimitEnforcer> limitEnforcer,
       kj::Own<void> ioContextDependency,
@@ -153,7 +153,7 @@ class WorkerEntrypoint::ResponseSentTracker final: public kj::HttpService::Respo
 };
 
 kj::Own<WorkerInterface> WorkerEntrypoint::construct(ThreadContext& threadContext,
-    kj::Own<const Worker> worker,
+    kj::Arc<Worker> worker,
     kj::Maybe<kj::StringPtr> entrypointName,
     Frankenvalue props,
     kj::Maybe<kj::Own<Worker::Actor>> actor,
@@ -197,7 +197,7 @@ WorkerEntrypoint::WorkerEntrypoint(kj::Badge<WorkerEntrypoint> badge,
       props(kj::mv(props)),
       cfBlobJson(kj::mv(cfBlobJson)) {}
 
-void WorkerEntrypoint::init(kj::Own<const Worker> worker,
+void WorkerEntrypoint::init(kj::Arc<Worker> worker,
     kj::Maybe<kj::Own<Worker::Actor>> actor,
     kj::Own<LimitEnforcer> limitEnforcer,
     kj::Own<void> ioContextDependency,
@@ -753,10 +753,10 @@ kj::Promise<WorkerInterface::CustomEvent::Result> WorkerEntrypoint::customEvent(
 }
 
 #ifdef KJ_DEBUG
-void requestGc(const Worker& worker) {
+void requestGc(kj::Arc<Worker> worker) {
   TRACE_EVENT("workerd", "Debug: requestGc()");
   jsg::runInV8Stack([&](jsg::V8StackScope& stackScope) {
-    auto& isolate = worker.getIsolate();
+    auto& isolate = worker->getIsolate();
     auto lock = isolate.getApi().lock(stackScope);
     lock->requestGcForTesting();
   });
@@ -765,13 +765,13 @@ void requestGc(const Worker& worker) {
 template <typename T>
 kj::Promise<T> addGcPassForTest(IoContext& context, kj::Promise<T> promise) {
   TRACE_EVENT("workerd", "Debug: addGcPassForTest");
-  auto worker = kj::atomicAddRef(context.getWorker());
+  auto worker = context.getWorker().addRef();
   if constexpr (kj::isSameType<T, void>()) {
     co_await promise;
-    requestGc(*worker);
+    requestGc(kj::mv(worker));
   } else {
     auto ret = co_await promise;
-    requestGc(*worker);
+    requestGc(kj::mv(worker));
     co_return kj::mv(ret);
   }
 }
@@ -790,7 +790,7 @@ kj::Promise<T> WorkerEntrypoint::maybeAddGcPassForTest(IoContext& context, kj::P
 }  // namespace
 
 kj::Own<WorkerInterface> newWorkerEntrypoint(ThreadContext& threadContext,
-    kj::Own<const Worker> worker,
+    kj::Arc<Worker> worker,
     kj::Maybe<kj::StringPtr> entrypointName,
     Frankenvalue props,
     kj::Maybe<kj::Own<Worker::Actor>> actor,

src/workerd/io/worker-entrypoint.h

@@ -29,7 +29,7 @@ class InvocationSpanContext;
 //   - Or, falling back to proxying if passThroughOnException() was used.
 // - Finish waitUntil() tasks.
 kj::Own<WorkerInterface> newWorkerEntrypoint(ThreadContext& threadContext,
-    kj::Own<const Worker> worker,
+    kj::Arc<Worker> worker,
     kj::Maybe<kj::StringPtr> entrypointName,
     Frankenvalue props,
     kj::Maybe<kj::Own<Worker::Actor>> actor,

src/workerd/io/worker.c++

@@ -803,7 +803,7 @@ struct Worker::Script::Impl {
     KJ_ASSERT(!FeatureFlags::get(js).getNewModuleRegistry(),
         "legacy dynamic imports must not be used with the new module registry");
     static auto constexpr handleDynamicImport =
-        [](kj::Own<const Worker> worker, DynamicImportHandler handler,
+        [](kj::Arc<Worker> worker, DynamicImportHandler handler,
             kj::Maybe<jsg::Ref<jsg::AsyncContextFrame>> asyncContext)
         -> kj::Promise<DynamicImportResult> {
       co_await kj::yield();
@@ -853,7 +853,7 @@ struct Worker::Script::Impl {
         auto& context = IoContext::current();
 
         return context.awaitIo(js,
-            handleDynamicImport(kj::atomicAddRef(context.getWorker()), kj::mv(handler),
+            handleDynamicImport(context.getWorker().addRef(), kj::mv(handler),
                 jsg::AsyncContextFrame::currentRef(js)),
             [](jsg::Lock& js, DynamicImportResult result) {
           if (result.isException) {
@@ -3431,7 +3431,7 @@ kj::Maybe<Worker::ConnectFn&> Worker::getConnectOverride(kj::StringPtr networkAd
   return connectOverrides.find(networkAddress);
 }
 
-Worker::Actor::Actor(const Worker& worker,
+Worker::Actor::Actor(kj::Arc<Worker> worker,
     kj::Maybe<RequestTracker&> tracker,
     Actor::Id actorId,
     bool hasTransient,
@@ -3445,7 +3445,7 @@ Worker::Actor::Actor(const Worker& worker,
     kj::Maybe<kj::Own<HibernationManager>> manager,
     kj::Maybe<uint16_t> hibernationEventType,
     kj::Maybe<rpc::Container::Client> container)
-    : worker(kj::atomicAddRef(worker)),
+    : worker(kj::mv(worker)),
       tracker(tracker.map([](RequestTracker& tracker) { return tracker.addRef(); })) {
   impl = kj::heap<Impl>(*this, lock, kj::mv(actorId), hasTransient, kj::mv(makeActorCache),
       kj::mv(makeStorage), kj::mv(loopback), timerChannel, kj::mv(metrics), kj::mv(manager),

src/workerd/io/worker.h

@@ -763,7 +763,7 @@ class Worker::Actor final: public kj::Refcounted {
 
   // Create a new Actor hosted by this Worker. Note that this Actor object may only be manipulated
   // from the thread that created it.
-  Actor(const Worker& worker,
+  Actor(kj::Arc<Worker> worker,
       kj::Maybe<RequestTracker&> tracker,
       Id actorId,
       bool hasTransient,
@@ -849,8 +849,8 @@ class Worker::Actor final: public kj::Refcounted {
   // Only needs to be called when allocating a HibernationManager!
   kj::Maybe<uint16_t> getHibernationEventType();
 
-  inline const Worker& getWorker() {
-    return *worker;
+  inline kj::Arc<Worker> getWorker() {
+    return worker.addRef();
   }
 
   void assertCanSetAlarm();
@@ -872,7 +872,7 @@ class Worker::Actor final: public kj::Refcounted {
  private:
   kj::Promise<WorkerInterface::ScheduleAlarmResult> handleAlarm(kj::Date scheduledTime);
 
-  kj::Own<const Worker> worker;
+  kj::Arc<Worker> worker;
   kj::Maybe<kj::Own<RequestTracker>> tracker;
   struct Impl;
   kj::Own<Impl> impl;

src/workerd/server/server.c++

@@ -1734,7 +1734,7 @@ class Server::WorkerService final: public Service,
   using AbortActorsCallback = kj::Function<void()>;
 
   WorkerService(ThreadContext& threadContext,
-      kj::Own<const Worker> worker,
+      kj::Arc<Worker> worker,
       kj::Maybe<kj::HashSet<kj::String>> defaultEntrypointHandlers,
       kj::HashMap<kj::String, kj::HashSet<kj::String>> namedEntrypoints,
       kj::HashSet<kj::String> actorClassEntrypoints,
@@ -1921,8 +1921,8 @@ class Server::WorkerService final: public Service,
     observer = kj::refcounted<RequestObserverWithTracer>(
         mapAddRef(workerTracer), kj::mv(streamingTailWorkers), waitUntilTasks);
 
-    return newWorkerEntrypoint(threadContext, kj::atomicAddRef(*worker), entrypointName,
-        kj::mv(props), kj::mv(actor), kj::Own<LimitEnforcer>(this, kj::NullDisposer::instance),
+    return newWorkerEntrypoint(threadContext, worker.addRef(), entrypointName, kj::mv(props),
+        kj::mv(actor), kj::Own<LimitEnforcer>(this, kj::NullDisposer::instance),
         {},  // ioContextDependency
         kj::Own<IoChannelFactory>(this, kj::NullDisposer::instance), kj::mv(observer),
         waitUntilTasks,
@@ -2081,16 +2081,15 @@ class Server::WorkerService final: public Service,
             co_return;
           }
           KJ_IF_SOME(m, manager) {
-            auto& worker = a->getWorker();
-            auto workerStrongRef = kj::atomicAddRef(worker);
+            auto worker = a->getWorker();
             // Take an async lock, we can't use `takeAsyncLock(RequestObserver&)` since we don't
             // have an `IncomingRequest` at this point.
             //
             // Note that we do not have a race here because this is part of the `shutdownTask`
             // promise. If a new request comes in while we're waiting to get the lock then we will
             // cancel this promise.
-            Worker::AsyncLock asyncLock = co_await worker.takeAsyncLockWithoutRequest(nullptr);
-            workerStrongRef->runInLockScope(
+            Worker::AsyncLock asyncLock = co_await worker->takeAsyncLockWithoutRequest(nullptr);
+            worker->runInLockScope(
                 asyncLock, [&](Worker::Lock& lock) { m->hibernateWebSockets(lock); });
           }
           a->shutdown(
@@ -2375,7 +2374,7 @@ class Server::WorkerService final: public Service,
             static constexpr uint16_t hibernationEventTypeId = 8;
 
             actorContainer->actor.emplace(
-                kj::refcounted<Worker::Actor>(*service.worker, actorContainer->getTracker(),
+                kj::refcounted<Worker::Actor>(service.worker.addRef(), actorContainer->getTracker(),
                     kj::str(idPtr), true, kj::mv(makeActorCache), className, kj::mv(makeStorage),
                     lock, kj::mv(loopback), timerChannel, kj::refcounted<ActorObserver>(),
                     actorContainer->tryGetManagerRef(), hibernationEventTypeId));
@@ -2466,7 +2465,7 @@ class Server::WorkerService final: public Service,
   // LinkedIoChannels owns the SqliteDatabase::Vfs, so make sure it is destroyed last.
   kj::OneOf<LinkCallback, LinkedIoChannels> ioChannels;
 
-  kj::Own<const Worker> worker;
+  kj::Arc<Worker> worker;
   kj::Maybe<kj::HashSet<kj::String>> defaultEntrypointHandlers;
   kj::HashMap<kj::String, kj::HashSet<kj::String>> namedEntrypoints;
   kj::HashSet<kj::String> actorClassEntrypoints;
@@ -3285,7 +3284,7 @@ kj::Own<Server::Service> Server::makeWorker(kj::StringPtr name,
   }
 
   jsg::V8Ref<v8::Object> ctxExportsHandle = nullptr;
-  auto worker = kj::atomicRefcounted<Worker>(kj::mv(script), kj::atomicRefcounted<WorkerObserver>(),
+  auto worker = kj::arc<Worker>(kj::mv(script), kj::atomicRefcounted<WorkerObserver>(),
       [&](jsg::Lock& lock, const Worker::Api& api, v8::Local<v8::Object> target,
           v8::Local<v8::Object> ctxExports) {
     // We can't fill in ctx.exports yet because we need to run the validator first to discover

src/workerd/tests/test-fixture.c++

@@ -343,7 +343,7 @@ TestFixture::TestFixture(SetupParams&& params)
           IsolateObserver::StartType::COLD,
           false,
           nullptr)),
-      worker(kj::atomicRefcounted<Worker>(kj::atomicAddRef(*workerScript),
+      worker(kj::arc<Worker>(kj::atomicAddRef(*workerScript),
           kj::atomicRefcounted<WorkerObserver>(),
           [](jsg::Lock&, const Worker::Api&, v8::Local<v8::Object>, v8::Local<v8::Object>) {
             // no bindings, nothing to do
@@ -367,7 +367,7 @@ TestFixture::TestFixture(SetupParams&& params)
         return jsg::alloc<api::DurableObjectStorage>(
             IoContext::current().addObject(actorCache), /*enableSql=*/false);
       };
-      actor = kj::refcounted<Worker::Actor>(*worker, /*tracker=*/kj::none, kj::mv(id),
+      actor = kj::refcounted<Worker::Actor>(worker.addRef(), /*tracker=*/kj::none, kj::mv(id),
           /*hasTransient=*/false, makeActorCache,
           /*classname=*/kj::none, makeStorage, lock, kj::refcounted<MockActorLoopback>(),
           *timerChannel, kj::refcounted<ActorObserver>(), kj::none, kj::none);
@@ -408,7 +408,7 @@ void TestFixture::runInIoContext(kj::Function<kj::Promise<void>(const Environmen
 
 kj::Own<IoContext::IncomingRequest> TestFixture::createIncomingRequest() {
   auto context = kj::refcounted<IoContext>(
-      threadContext, kj::atomicAddRef(*worker), actor, kj::heap<MockLimitEnforcer>());
+      threadContext, worker.addRef(), actor, kj::heap<MockLimitEnforcer>());
   auto invocationSpanContext = tracing::InvocationSpanContext::newForInvocation(kj::none, kj::none);
   auto incomingRequest = kj::heap<IoContext::IncomingRequest>(kj::addRef(*context),
       kj::heap<DummyIoChannelFactory>(*timerChannel), kj::refcounted<RequestObserver>(), nullptr,