Refactor Observer and LimitEnforcer cleaner for...

src/workerd/io/limit-enforcer.h

@@ -17,7 +17,7 @@ static constexpr size_t DEFAULT_MAX_PBKDF2_ITERATIONS = 100'000;
 // Interface for an object that enforces resource limits on an Isolate level.
 //
 // See also LimitEnforcer, which enforces on a per-request level.
-class IsolateLimitEnforcer {
+class IsolateLimitEnforcer: public kj::Refcounted {
 public:
   // Get CreateParams to pass when constructing a new isolate.
   virtual v8::Isolate::CreateParams getCreateParams() = 0;

src/workerd/io/observer.h

@@ -141,7 +141,9 @@ class RequestObserver: public kj::Refcounted {
   }
 };
 
-class IsolateObserver: public kj::AtomicRefcounted, public jsg::IsolateObserver {
+class JsgIsolateObserver: public kj::AtomicRefcounted, public jsg::IsolateObserver {};
+
+class IsolateObserver: public kj::AtomicRefcounted {
 public:
   virtual ~IsolateObserver() noexcept(false) {}
 

src/workerd/io/worker.c++

@@ -672,6 +672,8 @@ struct Worker::Isolate::Impl {
         actorCacheLru(limitEnforcer.getActorCacheLruOptions()) {
     jsg::runInV8Stack([&](jsg::V8StackScope& stackScope) {
       auto lock = api.lock(stackScope);
+
+      limitEnforcer.customizeIsolate(lock->v8Isolate);
       if (inspectorPolicy != InspectorPolicy::DISALLOW) {
         // We just created our isolate, so we don't need to use Isolate::Impl::Lock.
         KJ_ASSERT(!isMultiTenantProcess(), "inspector is not safe in multi-tenant processes");
@@ -969,18 +971,22 @@ const HeapSnapshotDeleter HeapSnapshotDeleter::INSTANCE;
 }  // namespace
 
 Worker::Isolate::Isolate(kj::Own<Api> apiParam,
+    kj::Own<IsolateObserver> metricsParam,
     kj::StringPtr id,
+    kj::Own<IsolateLimitEnforcer> limitEnforcerParam,
     InspectorPolicy inspectorPolicy,
     ConsoleMode consoleMode)
-    : metrics(kj::atomicAddRef(apiParam->getMetrics())),
+    : metrics(kj::mv(metricsParam)),
       id(kj::str(id)),
+      limitEnforcer(kj::mv(limitEnforcerParam)),
       api(kj::mv(apiParam)),
-      limitEnforcer(api->getLimitEnforcer()),
       consoleMode(consoleMode),
       featureFlagsForFl(makeCompatJson(decompileCompatibilityFlagsForFl(api->getFeatureFlags()))),
-      impl(kj::heap<Impl>(*api, *metrics, limitEnforcer, inspectorPolicy)),
+      impl(kj::heap<Impl>(*api, *metrics, *limitEnforcer, inspectorPolicy)),
       weakIsolateRef(WeakIsolateRef::wrap(this)),
       traceAsyncContextKey(kj::refcounted<jsg::AsyncContextFrame::StorageKey>()) {
+  api->setEnforcer(*limitEnforcer);
+  api->setIsolateObserver(*metrics);
   metrics->created();
   // We just created our isolate, so we don't need to use Isolate::Impl::Lock (nor an async lock).
   jsg::runInV8Stack([&](jsg::V8StackScope& stackScope) {
@@ -1325,7 +1331,7 @@ Worker::Script::Script(kj::Own<const Isolate> isolateParam,
             KJ_SWITCH_ONEOF(source) {
               KJ_CASE_ONEOF(script, ScriptSource) {
                 impl->globals =
-                    script.compileGlobals(lock, isolate->getApi(), isolate->impl->metrics);
+                    script.compileGlobals(lock, isolate->getApi(), isolate->getApi().getObserver());
 
                 {
                   // It's unclear to me if CompileUnboundScript() can get trapped in any infinite loops or
@@ -1378,7 +1384,7 @@ Worker::Isolate::~Isolate() noexcept(false) {
 
   // Update the isolate stats one last time to make sure we're accurate for cleanup in
   // `evicted()`.
-  limitEnforcer.reportMetrics(*metrics);
+  limitEnforcer->reportMetrics(*metrics);
 
   metrics->evicted();
   weakIsolateRef->invalidate();
@@ -1393,6 +1399,8 @@ Worker::Isolate::~Isolate() noexcept(false) {
     auto inspector = kj::mv(impl->inspector);
     auto dropTraceAsyncContextKey = kj::mv(traceAsyncContextKey);
   });
+  api->invalidateEnforcer();
+  api->invalidateIsolateObserver();
 }
 
 Worker::Script::~Script() noexcept(false) {
@@ -3793,7 +3801,7 @@ kj::Own<const Worker::Script> Worker::Isolate::newScript(kj::StringPtr scriptId,
 }
 
 void Worker::Isolate::completedRequest() const {
-  limitEnforcer.completedRequest(id);
+  limitEnforcer->completedRequest(id);
 }
 
 bool Worker::Isolate::isInspectorEnabled() const {

src/workerd/io/worker.h

@@ -275,7 +275,9 @@ class Worker::Isolate: public kj::AtomicRefcounted {
   // session has one isolate which may load many iterations of the script (this allows the
   // inspector session to stay open across them).
   explicit Isolate(kj::Own<Api> api,
+      kj::Own<IsolateObserver> metrics,
       kj::StringPtr id,
+      kj::Own<IsolateLimitEnforcer> limitEnforcer,
       InspectorPolicy inspectorPolicy,
       ConsoleMode consoleMode = ConsoleMode::INSPECTOR_ONLY);
 
@@ -305,11 +307,11 @@ class Worker::Isolate: public kj::AtomicRefcounted {
       kj::Maybe<ValidationErrorReporter&> errorReporter = kj::none) const;
 
   inline IsolateLimitEnforcer& getLimitEnforcer() {
-    return limitEnforcer;
+    return *limitEnforcer;
   }
 
   inline const IsolateLimitEnforcer& getLimitEnforcer() const {
-    return limitEnforcer;
+    return *limitEnforcer;
   }
 
   inline Api& getApi() {
@@ -408,8 +410,8 @@ class Worker::Isolate: public kj::AtomicRefcounted {
   TeardownFinishedGuard<IsolateObserver&> teardownGuard{*metrics};
 
   kj::String id;
+  kj::Own<IsolateLimitEnforcer> limitEnforcer;
   kj::Own<Api> api;
-  IsolateLimitEnforcer& limitEnforcer;
   ConsoleMode consoleMode;
 
   // If non-null, a serialized JSON object with a single "flags" property, which is a list of
@@ -532,10 +534,11 @@ class Worker::Api {
   virtual jsg::JsObject wrapExecutionContext(
       jsg::Lock& lock, jsg::Ref<api::ExecutionContext> ref) const = 0;
 
-  virtual IsolateLimitEnforcer& getLimitEnforcer() = 0;
-  virtual const IsolateLimitEnforcer& getLimitEnforcer() const = 0;
-  virtual IsolateObserver& getMetrics() = 0;
-  virtual const IsolateObserver& getMetrics() const = 0;
+  virtual const jsg::IsolateObserver& getObserver() const = 0;
+  virtual void setIsolateObserver(IsolateObserver&) = 0;
+  virtual void invalidateIsolateObserver() = 0;
+  virtual void setEnforcer(IsolateLimitEnforcer&) = 0;
+  virtual void invalidateEnforcer() = 0;
 
   // Set the module fallback service callback, if any.
   using ModuleFallbackCallback = kj::Maybe<kj::OneOf<kj::String, jsg::ModuleRegistry::ModuleInfo>>(

src/workerd/jsg/disable-memoize-test.c++

@@ -137,7 +137,7 @@ void expectEval(
 }
 
 KJ_TEST("Create a context with memoization disabled change flags then create another context") {
-  auto observer = kj::atomicRefcounted<workerd::IsolateObserver>();
+  auto observer = kj::atomicRefcounted<JsgIsolateObserver>();
   capnp::MallocMessageBuilder flagsArena;
   auto flags = flagsArena.initRoot<::workerd::CompatibilityFlags>();
   auto flagsReader = flags.asReader();

src/workerd/server/server.c++

@@ -2928,27 +2928,29 @@ kj::Own<Server::Service> Server::makeWorker(kj::StringPtr name,
     }
   };
 
+  auto jsgobserver = kj::atomicRefcounted<JsgIsolateObserver>();
   auto observer = kj::atomicRefcounted<IsolateObserver>();
-  auto limitEnforcer = kj::heap<NullIsolateLimitEnforcer>();
+  auto limitEnforcer = kj::refcounted<NullIsolateLimitEnforcer>();
 
   kj::Maybe<kj::Own<jsg::modules::ModuleRegistry>> newModuleRegistry;
   if (featureFlags.getNewModuleRegistry()) {
     KJ_REQUIRE(experimental,
         "The new ModuleRegistry implementation is an experimental feature. "
         "You must run workerd with `--experimental` to use this feature.");
     newModuleRegistry = WorkerdApi::initializeBundleModuleRegistry(
-        *observer, conf, featureFlags.asReader(), pythonConfig);
+        *jsgobserver, conf, featureFlags.asReader(), pythonConfig);
   }
 
-  auto api =
-      kj::heap<WorkerdApi>(globalContext->v8System, featureFlags.asReader(), kj::mv(limitEnforcer),
-          kj::mv(observer), *memoryCacheProvider, pythonConfig, kj::mv(newModuleRegistry));
+  auto api = kj::heap<WorkerdApi>(globalContext->v8System, featureFlags.asReader(),
+      limitEnforcer->getCreateParams(), kj::mv(jsgobserver), *memoryCacheProvider, pythonConfig,
+      kj::mv(newModuleRegistry));
   auto inspectorPolicy = Worker::Isolate::InspectorPolicy::DISALLOW;
   if (inspectorOverride != kj::none) {
     // For workerd, if the inspector is enabled, it is always fully trusted.
     inspectorPolicy = Worker::Isolate::InspectorPolicy::ALLOW_FULLY_TRUSTED;
   }
-  auto isolate = kj::atomicRefcounted<Worker::Isolate>(kj::mv(api), name, inspectorPolicy,
+  auto isolate = kj::atomicRefcounted<Worker::Isolate>(kj::mv(api), kj::mv(observer), name,
+      kj::mv(limitEnforcer), inspectorPolicy,
       conf.isServiceWorkerScript() ? Worker::ConsoleMode::INSPECTOR_ONLY : consoleMode);
 
   // If we are using the inspector, we need to register the Worker::Isolate

src/workerd/server/workerd-api.c++

@@ -131,8 +131,7 @@ static const PythonConfig defaultConfig{
 struct WorkerdApi::Impl final {
   kj::Own<CompatibilityFlags::Reader> features;
   kj::Maybe<kj::Own<jsg::modules::ModuleRegistry>> maybeOwnedModuleRegistry;
-  kj::Own<IsolateLimitEnforcer> limitEnforcer;
-  kj::Own<IsolateObserver> observer;
+  kj::Own<JsgIsolateObserver> observer;
   JsgWorkerdIsolate jsgIsolate;
   api::MemoryCacheProvider& memoryCacheProvider;
   const PythonConfig& pythonConfig;
@@ -159,24 +158,17 @@ struct WorkerdApi::Impl final {
 
   Impl(jsg::V8System& v8System,
       CompatibilityFlags::Reader featuresParam,
-      kj::Own<IsolateLimitEnforcer> limitEnforcerParam,
-      kj::Own<IsolateObserver> observerParam,
+      v8::Isolate::CreateParams createParams,
+      kj::Own<JsgIsolateObserver> observerParam,
       api::MemoryCacheProvider& memoryCacheProvider,
       const PythonConfig& pythonConfig = defaultConfig,
       kj::Maybe<kj::Own<jsg::modules::ModuleRegistry>> newModuleRegistry = kj::none)
       : features(capnp::clone(featuresParam)),
         maybeOwnedModuleRegistry(kj::mv(newModuleRegistry)),
-        limitEnforcer(kj::mv(limitEnforcerParam)),
         observer(kj::atomicAddRef(*observerParam)),
-        jsgIsolate(v8System,
-            Configuration(*this),
-            kj::mv(observerParam),
-            limitEnforcer->getCreateParams()),
+        jsgIsolate(v8System, Configuration(*this), kj::mv(observerParam), kj::mv(createParams)),
         memoryCacheProvider(memoryCacheProvider),
-        pythonConfig(pythonConfig) {
-    jsgIsolate.runInLockScope(
-        [&](JsgWorkerdIsolate::Lock& lock) { limitEnforcer->customizeIsolate(lock.v8Isolate); });
-  }
+        pythonConfig(pythonConfig) {}
 
   static v8::Local<v8::String> compileTextGlobal(
       JsgWorkerdIsolate::Lock& lock, capnp::Text::Reader reader) {
@@ -218,14 +210,14 @@ struct WorkerdApi::Impl final {
 
 WorkerdApi::WorkerdApi(jsg::V8System& v8System,
     CompatibilityFlags::Reader features,
-    kj::Own<IsolateLimitEnforcer> limitEnforcer,
-    kj::Own<IsolateObserver> observer,
+    v8::Isolate::CreateParams createParams,
+    kj::Own<JsgIsolateObserver> observer,
     api::MemoryCacheProvider& memoryCacheProvider,
     const PythonConfig& pythonConfig,
     kj::Maybe<kj::Own<jsg::modules::ModuleRegistry>> newModuleRegistry)
     : impl(kj::heap<Impl>(v8System,
           features,
-          kj::mv(limitEnforcer),
+          kj::mv(createParams),
           kj::mv(observer),
           memoryCacheProvider,
           pythonConfig,
@@ -275,21 +267,14 @@ jsg::JsObject WorkerdApi::wrapExecutionContext(
       kj::downcast<JsgWorkerdIsolate::Lock>(lock).wrap(lock.v8Context(), kj::mv(ref)));
 }
 
-IsolateLimitEnforcer& WorkerdApi::getLimitEnforcer() {
-  return *impl->limitEnforcer;
-}
-
-const IsolateLimitEnforcer& WorkerdApi::getLimitEnforcer() const {
-  return *impl->limitEnforcer;
-}
-
-IsolateObserver& WorkerdApi::getMetrics() {
+const jsg::IsolateObserver& WorkerdApi::getObserver() const {
   return *impl->observer;
 }
 
-const IsolateObserver& WorkerdApi::getMetrics() const {
-  return *impl->observer;
-}
+void WorkerdApi::setIsolateObserver(IsolateObserver&) {};
+void WorkerdApi::invalidateIsolateObserver() {};
+void WorkerdApi::setEnforcer(IsolateLimitEnforcer&) {}
+void WorkerdApi::invalidateEnforcer() {}
 
 Worker::Script::Source WorkerdApi::extractSource(kj::StringPtr name,
     config::Worker::Reader conf,

src/workerd/server/workerd-api.h

@@ -34,8 +34,8 @@ class WorkerdApi final: public Worker::Api {
 public:
   WorkerdApi(jsg::V8System& v8System,
       CompatibilityFlags::Reader features,
-      kj::Own<IsolateLimitEnforcer> limitEnforcer,
-      kj::Own<IsolateObserver> observer,
+      v8::Isolate::CreateParams createParams,
+      kj::Own<JsgIsolateObserver> observer,
       api::MemoryCacheProvider& memoryCacheProvider,
       const PythonConfig& pythonConfig,
       kj::Maybe<kj::Own<jsg::modules::ModuleRegistry>> newModuleRegistry);
@@ -55,10 +55,11 @@ class WorkerdApi final: public Worker::Api {
       jsg::Lock& lock) const override;
   jsg::JsObject wrapExecutionContext(
       jsg::Lock& lock, jsg::Ref<api::ExecutionContext> ref) const override;
-  IsolateLimitEnforcer& getLimitEnforcer() override;
-  const IsolateLimitEnforcer& getLimitEnforcer() const override;
-  IsolateObserver& getMetrics() override;
-  const IsolateObserver& getMetrics() const override;
+  const jsg::IsolateObserver& getObserver() const override;
+  void setIsolateObserver(IsolateObserver&) override;
+  void invalidateIsolateObserver() override;
+  void setEnforcer(IsolateLimitEnforcer&) override;
+  void invalidateEnforcer() override;
 
   static Worker::Script::Source extractSource(kj::StringPtr name,
       config::Worker::Reader conf,

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

@@ -323,13 +323,16 @@ TestFixture::TestFixture(SetupParams&& params)
       memoryCacheProvider(kj::heap<api::MemoryCacheProvider>(*timer)),
       api(kj::heap<server::WorkerdApi>(testV8System,
           params.featureFlags.orDefault(CompatibilityFlags::Reader()),
-          kj::heap<MockIsolateLimitEnforcer>(),
-          kj::atomicRefcounted<IsolateObserver>(),
+          kj::heap<MockIsolateLimitEnforcer>()->getCreateParams(),
+          kj::atomicRefcounted<JsgIsolateObserver>(),
           *memoryCacheProvider,
           defaultPythonConfig,
           kj::none)),
-      workerIsolate(kj::atomicRefcounted<Worker::Isolate>(
-          kj::mv(api), scriptId, Worker::Isolate::InspectorPolicy::DISALLOW)),
+      workerIsolate(kj::atomicRefcounted<Worker::Isolate>(kj::mv(api),
+          kj::atomicRefcounted<IsolateObserver>(),
+          scriptId,
+          kj::heap<MockIsolateLimitEnforcer>(),
+          Worker::Isolate::InspectorPolicy::DISALLOW)),
       workerScript(kj::atomicRefcounted<Worker::Script>(kj::atomicAddRef(*workerIsolate),
           scriptId,
           server::WorkerdApi::extractSource(mainModuleName,