Refactor limitEnforcer from api to Isolate.

This will return ownership over limitEnforcer from the api class to the
isolate class.

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/worker.c++

@@ -663,7 +663,7 @@ struct Worker::Isolate::Impl {
   //   because our GlobalScope object needs to have a function called on it, and any attached
   //   inspector needs to be notified. JSG doesn't know about these things.
 
-  Impl(const Api& api,
+  Impl(Api& api,
       IsolateObserver& metrics,
       IsolateLimitEnforcer& limitEnforcer,
       InspectorPolicy inspectorPolicy)
@@ -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");
@@ -970,17 +972,19 @@ const HeapSnapshotDeleter HeapSnapshotDeleter::INSTANCE;
 
 Worker::Isolate::Isolate(kj::Own<Api> apiParam,
     kj::StringPtr id,
+    kj::Own<IsolateLimitEnforcer> limitEnforcerParam,
     InspectorPolicy inspectorPolicy,
     ConsoleMode consoleMode)
     : metrics(kj::atomicAddRef(apiParam->getMetrics())),
       id(kj::str(id)),
       api(kj::mv(apiParam)),
-      limitEnforcer(api->getLimitEnforcer()),
+      limitEnforcer(kj::mv(limitEnforcerParam)),
       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);
   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) {
@@ -1383,7 +1387,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();
@@ -1398,6 +1402,7 @@ Worker::Isolate::~Isolate() noexcept(false) {
     auto inspector = kj::mv(impl->inspector);
     auto dropTraceAsyncContextKey = kj::mv(traceAsyncContextKey);
   });
+  api->invalidateEnforcer();
 }
 
 Worker::Script::~Script() noexcept(false) {
@@ -3800,7 +3805,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

@@ -276,6 +276,7 @@ class Worker::Isolate: public kj::AtomicRefcounted {
   // inspector session to stay open across them).
   explicit Isolate(kj::Own<Api> api,
       kj::StringPtr id,
+      kj::Own<IsolateLimitEnforcer> limitEnforcer,
       InspectorPolicy inspectorPolicy,
       ConsoleMode consoleMode = ConsoleMode::INSPECTOR_ONLY);
 
@@ -305,11 +306,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() {
@@ -409,7 +410,8 @@ class Worker::Isolate: public kj::AtomicRefcounted {
 
   kj::String id;
   kj::Own<Api> api;
-  IsolateLimitEnforcer& limitEnforcer;
+  // TODO: should this be before or after api?
+  kj::Own<IsolateLimitEnforcer> limitEnforcer;
   ConsoleMode consoleMode;
 
   // If non-null, a serialized JSON object with a single "flags" property, which is a list of
@@ -532,10 +534,10 @@ 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 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/server/server.c++

@@ -2922,7 +2922,7 @@ kj::Own<Server::Service> Server::makeWorker(kj::StringPtr name,
   };
 
   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()) {
@@ -2933,15 +2933,16 @@ kj::Own<Server::Service> Server::makeWorker(kj::StringPtr name,
         *observer, 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(observer), *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), 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++

@@ -132,7 +132,6 @@ 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;
   JsgWorkerdIsolate jsgIsolate;
   api::MemoryCacheProvider& memoryCacheProvider;
@@ -160,19 +159,15 @@ struct WorkerdApi::Impl final {
 
   Impl(jsg::V8System& v8System,
       CompatibilityFlags::Reader featuresParam,
-      kj::Own<IsolateLimitEnforcer> limitEnforcerParam,
+      v8::Isolate::CreateParams createParams,
       kj::Own<IsolateObserver> 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(
@@ -219,14 +214,14 @@ struct WorkerdApi::Impl final {
 
 WorkerdApi::WorkerdApi(jsg::V8System& v8System,
     CompatibilityFlags::Reader features,
-    kj::Own<IsolateLimitEnforcer> limitEnforcer,
+    v8::Isolate::CreateParams createParams,
     kj::Own<IsolateObserver> 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,
@@ -276,14 +271,6 @@ 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() {
   return *impl->observer;
 }
@@ -292,6 +279,9 @@ const IsolateObserver& WorkerdApi::getMetrics() const {
   return *impl->observer;
 }
 
+void WorkerdApi::setEnforcer(IsolateLimitEnforcer&) {}
+void WorkerdApi::invalidateEnforcer() {}
+
 Worker::Script::Source WorkerdApi::extractSource(kj::StringPtr name,
     config::Worker::Reader conf,
     Worker::ValidationErrorReporter& errorReporter,

src/workerd/server/workerd-api.h

@@ -34,7 +34,7 @@ class WorkerdApi final: public Worker::Api {
  public:
   WorkerdApi(jsg::V8System& v8System,
       CompatibilityFlags::Reader features,
-      kj::Own<IsolateLimitEnforcer> limitEnforcer,
+      v8::Isolate::CreateParams createParams,
       kj::Own<IsolateObserver> observer,
       api::MemoryCacheProvider& memoryCacheProvider,
       const PythonConfig& pythonConfig,
@@ -55,10 +55,10 @@ 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;
+  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,15 @@ 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::heap<MockIsolateLimitEnforcer>()->getCreateParams(),
           kj::atomicRefcounted<IsolateObserver>(),
           *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),
+          scriptId,
+          kj::heap<MockIsolateLimitEnforcer>(),
+          Worker::Isolate::InspectorPolicy::DISALLOW)),
       workerScript(kj::atomicRefcounted<Worker::Script>(kj::atomicAddRef(*workerIsolate),
           scriptId,
           server::WorkerdApi::extractSource(mainModuleName,